Cl 834

[muskansh-google]

logs - https://cloudlogging.app.goo.gl/bTerLAqLt2utW6Ky6

[gemini-code-assist]

Code Review

This pull request introduces support for Qwen3-Next models, including context-parallel evaluation of GatedDeltaNet inter-chunk recurrence, hybrid GDN implementations using Tokamax and Pallas, and hierarchical nested scans for scanned blocks. It also updates the Muon optimizer to handle scanned layers and adds associated unit tests. Key feedback includes parameterizing hardcoded user paths and removing a sleep command in the XPK script, addressing duplicate logical axis rules for 'embed', restoring TPU platform checks in ragged gather kernels to prevent crashes on non-TPU environments, and handling potential non-divisible cases in MoE group size calculations.

[gemini-code-assist]

This script contains several hardcoded user-specific paths and resource names, which makes it non-portable and difficult for other developers to use.

For example:

• Line 5: source /usr/local/google/home/muskansh/maxtext_env/bin/activate
• Line 30: cd /usr/local/google/home/muskansh/maxtext
• Line 180: client.bucket("muskansh-data")
• Line 186: PYTHONPATH=/usr/local/google/home/muskansh/xpk/src

These should be parameterized, for instance, by using environment variables defined at the top of the script or passed as arguments. This would greatly improve the script's reusability.

Suggestion:

At the top of the script, define variables for these paths and names, allowing them to be overridden by environment variables if they are already set:

# --- User-configurable paths and names ---
MAXTEXT_DIR="${MAXTEXT_DIR:-/usr/local/google/home/muskansh/maxtext}"
VENV_PATH="${VENV_PATH:-/usr/local/google/home/muskansh/maxtext_env}"
XPK_DIR="${XPK_DIR:-/usr/local/google/home/muskansh/xpk}"
GCS_BUCKET="${GCS_BUCKET:-muskansh-data}"

Then use these variables throughout the script, for example:

# line 5
source "${VENV_PATH}/bin/activate"

# line 30
cd "${MAXTEXT_DIR}" && \

# line 180
... client.bucket(\"${GCS_BUCKET}\") ...

# line 186
PYTHONPATH="${XPK_DIR}/src" python3 -P -m xpk.main workload create \

[gemini-code-assist]

The RUN_COMMAND ends with sleep 3600. This will keep the container running for an hour after the main training command finishes. While this can be useful for debugging, it's not ideal for a script in a shared repository as it can lead to resource wastage if not intended.

It would be better to make this behavior optional or configurable. For example, you could add an environment variable KEEP_ALIVE_SECS that defaults to 0.

Suggested change

	sleep 3600"
	sleep 0"

[gemini-code-assist]

There appear to be duplicate logical axis rules for 'embed'. Depending on how this configuration is parsed, the first rule might be overridden by the second, or it could lead to unexpected behavior. It's best to remove the redundant or incorrect entry to avoid confusion and potential issues.

[gemini-code-assist]

The platform check jax.devices()[0].platform != "tpu" has been removed. The following line sc_info = pltpu.get_tpu_info().sparse_core will crash on non-TPU platforms if enforce_fallback is not explicitly set to True. This change makes the function less safe to call directly on non-TPU environments and shifts the responsibility of platform checking to the caller. Consider restoring the platform check to prevent potential runtime errors.

Suggested change

	if enforce_fallback:
	if enforce_fallback or jax.devices()[0].platform != "tpu":

[gemini-code-assist]

The platform check jax.devices()[0].platform != "tpu" has been removed. The following line sc_info = pltpu.get_tpu_info().sparse_core will crash on non-TPU platforms if enforce_fallback is not explicitly set to True. This change makes the function less safe to call directly on non-TPU environments and shifts the responsibility of platform checking to the caller. Consider restoring the platform check to prevent potential runtime errors.

Suggested change

	if enforce_fallback:
	if enforce_fallback or jax.devices()[0].platform != "tpu":

[gemini-code-assist]

The calculation of avg_size uses integer division: avg_size = inputs.shape[0] // num_groups. If inputs.shape[0] is not perfectly divisible by num_groups, the sum of the sizes in the resulting tuple (avg_size,) * num_groups will not equal inputs.shape[0]. This could lead to incorrect behavior or errors in tokamax.RaggedDotGroupSizes. Please ensure this is the intended behavior or add handling for the non-divisible case.