Cl 974697834 - #5127
Conversation
…uon sharding, and mHC-lite Pallas/Mosaic kernels on TPU v6e-256. - Integrate DeepSeek mHC-lite architecture and high-performance Mosaic/Pallas TPU kernels with custom VJP rules and SPMD shard_map integration. - Support Qwen3-Next 80B architecture with Gated Delta Net (GDNv3) and Pallas custom VJP in hybrid_gdn.py. - Configure Muon optimizer and sharding rules for MoE parameters (shard-exp-on-fsdp). - Update AOT compile, XManager Borg, and XPK launch scripts for Qwen3-Next 80B on TPU v6e-256 (glp_16x16) topology to enable mHC-lite and Pallas kernel flags. - Add and update unit tests for mHC, GDN, Muon utilities, NNX decoders, and parameter mapping. - Validated clean AOT compilation on TPU v6e-256 via qn80b-v6e256-compile.sh and unit tests via mhc_test.
- Integrates sharding-aware Muon optimizer supporting Newton-Schulz iterations and partition padding into MaxText. - Fixes weight dimension indexing in muon_utils.py to account for layer scanning, 1D tensor exclusions, and rank-aware 2D attention projection handling. - Fixes routing gate exclusion logic to use exact segment matching across gate and router variants. - Plumbs NamedSharding in sharding.py by collecting momentum trees across multi-branch partitioned optimizers, extracting MuonState NamedTuples, and filtering optax.MaskedNode placeholders. - Threads mesh through train_utils.py, train_compile.py, and maxtext_engine.py to get_optimizer. - Add some minor tests throughout.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for the Qwen3-Next model architecture, adds experimental SFT and evaluation pipelines for the Omni (Gemma 3 Vision + Qwen 3 LLM) model, transitions checkpointing to the Orbax v1 API, and implements various kernel and optimizer enhancements such as hybrid GatedDeltaNet (GDN) and Muon improvements. The code review identified several critical runtime bugs, including invalid JAX API usage (such as jax.typeof and an unsupported mode argument in .at().set()), missing epath imports in GCS utilities, potential IndexError and KeyError exceptions in utility functions, and a region mismatch in the XPK runner script.
| # Scatter the per-slot dot back to flat slot order; dropped slots stay zero. | ||
| dot_local = jnp.sum(gathered.astype(jnp.float32) * sorted_tokens_local.astype(jnp.float32), axis=-1) | ||
| slots = jnp.where(mask, sliced_idx_inv, n) | ||
| grad_topk_weights = jnp.zeros((n,), jnp.float32).at[slots].set(dot_local, mode="drop") |
There was a problem hiding this comment.
The mode argument is not supported by JAX's .at[...].set(...) method and will raise a TypeError at runtime. Since JAX automatically ignores/drops out-of-bounds indices by default, the mode="drop" argument is both invalid and unnecessary.
| grad_topk_weights = jnp.zeros((n,), jnp.float32).at[slots].set(dot_local, mode="drop") | |
| grad_topk_weights = jnp.zeros((n,), jnp.float32).at[slots].set(dot_local) |
| activation_spec = jax.typeof(y).sharding.spec | ||
| scale_spec = jax.typeof(scale).sharding.spec |
There was a problem hiding this comment.
jax.typeof is not a valid JAX API and will raise an AttributeError at runtime. You can directly access the .sharding property of JAX arrays (e.g., y.sharding.spec and scale.sharding.spec).
| activation_spec = jax.typeof(y).sharding.spec | |
| scale_spec = jax.typeof(scale).sharding.spec | |
| activation_spec = y.sharding.spec | |
| scale_spec = scale.sharding.spec |
| while isinstance(hf_source_keys, list): | ||
| depth += 1 | ||
| hf_source_keys = hf_source_keys[0] |
There was a problem hiding this comment.
If hf_source_keys is an empty list [], accessing hf_source_keys[0] will raise an IndexError: list index out of range. We should defensively check that the list is not empty before accessing its first element.
| while isinstance(hf_source_keys, list): | |
| depth += 1 | |
| hf_source_keys = hf_source_keys[0] | |
| while isinstance(hf_source_keys, list) and hf_source_keys: | |
| depth += 1 | |
| hf_source_keys = hf_source_keys[0] |
| if not destination_gcs_name.startswith("gs://"): | ||
| dest = epath.Path(destination_gcs_name) | ||
| dest.parent.mkdir(parents=True, exist_ok=True) | ||
| epath.Path(source_file_name).copy(dest, overwrite=True) | ||
| return |
There was a problem hiding this comment.
epath is not imported in gcs_utils.py, which will lead to a NameError at runtime. Since gcs_utils.py already imports os and shutil, we can use standard library functions os.makedirs and shutil.copy to perform these local filesystem operations safely and cleanly.
| if not destination_gcs_name.startswith("gs://"): | |
| dest = epath.Path(destination_gcs_name) | |
| dest.parent.mkdir(parents=True, exist_ok=True) | |
| epath.Path(source_file_name).copy(dest, overwrite=True) | |
| return | |
| if not destination_gcs_name.startswith("gs://"): | |
| os.makedirs(os.path.dirname(destination_gcs_name), exist_ok=True) | |
| shutil.copy(source_file_name, destination_gcs_name) | |
| return |
| if not target_dir.startswith("gs://"): | ||
| dest_dir = epath.Path(target_dir) | ||
| dest_dir.mkdir(parents=True, exist_ok=True) | ||
| for root, _, files in os.walk(local_dir): | ||
| for file in files: | ||
| if module_name and module_name not in file: | ||
| continue | ||
| else: | ||
| max_logging.log(f"Uploading {file}") | ||
| local_path = os.path.join(root, file) | ||
| relative_path = os.path.relpath(local_path, local_dir) | ||
| dest_path = dest_dir / relative_path | ||
| dest_path.parent.mkdir(parents=True, exist_ok=True) | ||
| epath.Path(local_path).copy(dest_path, overwrite=True) |
There was a problem hiding this comment.
epath is not imported in gcs_utils.py, which will lead to a NameError at runtime. Since gcs_utils.py already imports os and shutil, we can use standard library functions os.makedirs and shutil.copy to perform these local filesystem operations safely and cleanly.
| if not target_dir.startswith("gs://"): | |
| dest_dir = epath.Path(target_dir) | |
| dest_dir.mkdir(parents=True, exist_ok=True) | |
| for root, _, files in os.walk(local_dir): | |
| for file in files: | |
| if module_name and module_name not in file: | |
| continue | |
| else: | |
| max_logging.log(f"Uploading {file}") | |
| local_path = os.path.join(root, file) | |
| relative_path = os.path.relpath(local_path, local_dir) | |
| dest_path = dest_dir / relative_path | |
| dest_path.parent.mkdir(parents=True, exist_ok=True) | |
| epath.Path(local_path).copy(dest_path, overwrite=True) | |
| if not target_dir.startswith("gs://"): | |
| os.makedirs(target_dir, exist_ok=True) | |
| for root, _, files in os.walk(local_dir): | |
| for file in files: | |
| if module_name and module_name not in file: | |
| continue | |
| else: | |
| max_logging.log(f"Uploading {file}") | |
| local_path = os.path.join(root, file) | |
| relative_path = os.path.relpath(local_path, local_dir) | |
| dest_path = os.path.join(target_dir, relative_path) | |
| os.makedirs(os.path.dirname(dest_path), exist_ok=True) | |
| shutil.copy(local_path, dest_path) |
| LOGS_URL="https://console.cloud.google.com/logs/query;query=resource.type%3D%22k8s_container%22%0Aresource.labels.project_id%3D%22${PROJECT_ID}%22%0Aresource.labels.location%3D%22southamerica-west1%22%0Aresource.labels.cluster_name%3D%22${CLUSTER_NAME}%22%0Aresource.labels.namespace_name%3D%22default%22%0Aresource.labels.pod_name%3A%22${WORKLOAD_NAME}-slice-job-0-0-%22%0Aseverity%3E%3DDEFAULT;storageScope=project;duration=P1D?project=${PROJECT_ID}" | ||
| GKE_URL="https://console.cloud.google.com/kubernetes/service/southamerica-west1/${CLUSTER_NAME}/default/${WORKLOAD_NAME}/details?project=${PROJECT_ID}" |
There was a problem hiding this comment.
The region southamerica-west1 is hardcoded in both LOGS_URL and GKE_URL, but the script configures ZONE="us-central2". This discrepancy will generate incorrect links. We should use the ${ZONE} variable dynamically in the URLs instead of hardcoding a different region.
| LOGS_URL="https://console.cloud.google.com/logs/query;query=resource.type%3D%22k8s_container%22%0Aresource.labels.project_id%3D%22${PROJECT_ID}%22%0Aresource.labels.location%3D%22southamerica-west1%22%0Aresource.labels.cluster_name%3D%22${CLUSTER_NAME}%22%0Aresource.labels.namespace_name%3D%22default%22%0Aresource.labels.pod_name%3A%22${WORKLOAD_NAME}-slice-job-0-0-%22%0Aseverity%3E%3DDEFAULT;storageScope=project;duration=P1D?project=${PROJECT_ID}" | |
| GKE_URL="https://console.cloud.google.com/kubernetes/service/southamerica-west1/${CLUSTER_NAME}/default/${WORKLOAD_NAME}/details?project=${PROJECT_ID}" | |
| LOGS_URL="https://console.cloud.google.com/logs/query;query=resource.type%3D%22k8s_container%22%0Aresource.labels.project_id%3D%22${PROJECT_ID}%22%0Aresource.labels.location%3D%22${ZONE}%22%0Aresource.labels.cluster_name%3D%22${CLUSTER_NAME}%22%0Aresource.labels.namespace_name%3D%22default%22%0Aresource.labels.pod_name%3A%22${WORKLOAD_NAME}-slice-job-0-0-%22%0Aseverity%3E%3DDEFAULT;storageScope=project;duration=P1D?project=${PROJECT_ID}" | |
| GKE_URL="https://console.cloud.google.com/kubernetes/service/${ZONE}/${CLUSTER_NAME}/default/${WORKLOAD_NAME}/details?project=${PROJECT_ID}" |
| def _batch_axis_size(self, axis_names: Any) -> int: | ||
| """Product of mesh dims backing one PartitionSpec entry (None -> 1).""" | ||
| if axis_names is None: | ||
| return 1 | ||
| names = axis_names if isinstance(axis_names, tuple) else (axis_names,) | ||
| size = 1 | ||
| for name in names: | ||
| size *= self._mesh.shape[name] | ||
| return size |
There was a problem hiding this comment.
If name is not present in self._mesh.shape (e.g., due to a mismatch or custom mesh rules), accessing self._mesh.shape[name] will raise a KeyError. Using .get(name, 1) is safer and prevents potential runtime errors.
| def _batch_axis_size(self, axis_names: Any) -> int: | |
| """Product of mesh dims backing one PartitionSpec entry (None -> 1).""" | |
| if axis_names is None: | |
| return 1 | |
| names = axis_names if isinstance(axis_names, tuple) else (axis_names,) | |
| size = 1 | |
| for name in names: | |
| size *= self._mesh.shape[name] | |
| return size | |
| def _batch_axis_size(self, axis_names: Any) -> int: | |
| """Product of mesh dims backing one PartitionSpec entry (None -> 1).""" | |
| if axis_names is None: | |
| return 1 | |
| names = axis_names if isinstance(axis_names, tuple) else (axis_names,) | |
| size = 1 | |
| for name in names: | |
| size *= self._mesh.shape.get(name, 1) | |
| return size |
2 participants
Description
Start with a short description of what the PR does and how this is a change from
the past.
The rest of the description includes relevant details and context, examples:
If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456
You can also provide a comma-separated list. If you don't want to close a bug but
simply to reference it, use BUGS, e.g.:
BUGS: b/123456
Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.
Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.
Tests
Please describe how you tested this change, and include any instructions and/or
commands to reproduce.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.