Onboard Qwen3.5 397b fp8 by Shuwen-Fang · Pull Request #5182 · AI-Hypercomputer/maxtext

Shuwen-Fang · 2026-09-09T23:07:39Z

Description

Onboards Qwen/Qwen3.5-397B-A17B-FP8 to MaxText and fixes scale-tensor sharding for block-scaled FP8 modules:

Qwen3.5-397B-FP8 Model Onboarding:
- Adds qwen3.5-397b-a17b-fp8.yml configuring FP8 weight storage (weight_dtype: "float8_e4m3fn"), 128x128 block scaling (weight_block_size: 128), and full-precision preservation for unquantized_modules.
- Registers the model across HuggingFace checkpoint conversion and metadata mappings (hf_model_configs.py, hf_shape.py, param_mapping.py, globals.py, types.py).
Block Scale Sharding Fixes:
- DenseGeneral: Replicates block-compressed scale dimensions instead of inheriting the kernel's mesh axes. Inheriting kernel sharding causes divisibility crashes when FSDP/TP mesh degree exceeds the compressed dimension (e.g., FSDP > 32 on a 32-block axis).
- RoutedMoE: Shards expert scales strictly along the expert axis (wi_kernel_axes[0]) and replicates the block-tile dimensions.
- Updates tests/unit/linears_test.py to assert scale replication on block-tiled layers.

Tests

JAX_PLATFORMS=cpu PYTHONPATH=src pytest tests/unit/configs_test.py tests/unit/pyconfig_test.py -q
JAX_PLATFORMS=cpu PYTHONPATH=src pytest tests/unit/linears_test.py -q
AOT train compilation verified for qwen3.5-397b-a17b-fp8

BUGS: b/557372531

Checklist

I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run end-to-end tests and provided workload links above if applicable.
I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

… and types

… llama 3.1 and qwen 3.5 fp8 - linears.py: simplify dequantize_weight to handle scalar (Llama 3.1) and block scales (Qwen 3.5); remove DeepSeekV4GroupedLinear quantization hooks and redundant parameter accessors. - common_types.py: replace cascading fnmatch loops in get_weight_dtype with static leaf lookup UNQUANTIZED_MODULE_LEAF_NAMES. - initializers.py: use explicit is_fp8_dtype instead of itemsize heuristics. - moe.py: replace np.ceil with integer division and clean up repeated hasattr guards. - maxengine.py: revert out-of-scope cache casting changes.

…ck-scaled quantization

…end_dtype

…ostic

…y to config

…stion in to_maxtext

…ides in to_maxtext

…pping

gemini-code-assist

Code Review

This pull request introduces support for FP8 weight-only storage with dynamic dequantization in MaxText, specifically targeting Qwen3.5 models. It implements dynamic dequantization utilities, integrates them into linear and MoE layers, and adds support for dequantizing weights on-the-fly during checkpoint loading. The review feedback highlights several critical issues, including a reversed loop structure in MoE expert parameter mapping that causes crashes, a runtime AttributeError in Qwen3NextRMSNorm from calling get_value() on an NNX parameter, and potential shared memory exhaustion from automatically setting HF_HOME to /dev/shm. Additionally, the reviewer recommends addressing code duplication for key resolution and tensor stacking, fixing a potential NameError with str2bool, and avoiding hardcoded absolute paths in test scripts.

gemini-code-assist · 2026-09-09T23:10:55Z

+if "HF_HOME" not in os.environ and os.path.exists("/dev/shm"):
+  os.environ["HF_HOME"] = "/dev/shm/hf_cache"


Setting HF_HOME to /dev/shm/hf_cache automatically when /dev/shm exists can easily exhaust the shared memory (shm) space, especially for extremely large models like Qwen3.5 397B (which is hundreds of gigabytes). Shared memory is typically limited in size and storing large model checkpoints there can lead to 'No space left on device' or bus errors. It is safer to rely on the default Hugging Face cache directory or let the user configure HF_HOME explicitly.

gemini-code-assist · 2026-09-09T23:10:55Z

+        y = jax.lax.with_sharding_constraint(y, out_sharding)
+      return y
+
+    scale = self.scale.get_value()


The method get_value() does not exist on flax.nnx.Param (or Variable). Calling it will raise an AttributeError at runtime. To get the parameter value, use self.scale.value or self.scale[...], which is consistent with how other parameters are accessed in this codebase (e.g., in linears.py and moe.py).

Suggested change

scale = self.scale.get_value()

scale = self.scale[...]

gemini-code-assist · 2026-09-09T23:10:55Z

+    resolved_key = key
+    shard_name = self.shard_map.get(resolved_key)
+    if shard_name is None:
+      # Check fallback for .weight_scale vs .scale and inverse scales
+      if resolved_key.endswith(".weight_scale"):
+        for suffix in [".scale", ".weight_scale_inv", ".scale_inv"]:
+          alt_key = resolved_key[:-len(".weight_scale")] + suffix
+          if alt_key in self.shard_map:
+            resolved_key = alt_key
+            shard_name = self.shard_map[resolved_key]
+            break
+      elif resolved_key.endswith(".scale"):
+        for suffix in [".weight_scale", ".scale_inv", ".weight_scale_inv"]:
+          alt_key = resolved_key[:-len(".scale")] + suffix
+          if alt_key in self.shard_map:
+            resolved_key = alt_key
+            shard_name = self.shard_map[resolved_key]
+            break
+


The fallback key resolution logic is duplicated here and inside the _ram_lock block. Extracting this into a local helper function _resolve_key makes the code cleaner, more maintainable, and eliminates the duplication.

def _resolve_key(k, container): if k in container: return k if k.endswith(".weight_scale"): base, suffixes = k[:-13], [".scale", ".weight_scale_inv", ".scale_inv"] elif k.endswith(".scale"): base, suffixes = k[:-6], [".weight_scale", ".scale_inv", ".weight_scale_inv"] else: return k for suffix in suffixes: alt_key = base + suffix if alt_key in container: return alt_key return k resolved_key = _resolve_key(key, self.shard_map) shard_name = self.shard_map.get(resolved_key)

gemini-code-assist · 2026-09-09T23:10:55Z

+        final_key = resolved_key
+        if final_key not in f.keys():
+          if final_key.endswith(".weight_scale"):
+            for suffix in [".scale", ".weight_scale_inv", ".scale_inv"]:
+              alt_key = final_key[:-len(".weight_scale")] + suffix
+              if alt_key in f.keys():
+                final_key = alt_key
+                break
+          elif final_key.endswith(".scale"):
+            for suffix in [".weight_scale", ".scale_inv", ".weight_scale_inv"]:
+              alt_key = final_key[:-len(".scale")] + suffix
+              if alt_key in f.keys():
+                final_key = alt_key
+                break
+        t = f.get_tensor(final_key)


Use the extracted _resolve_key helper function here to simplify the fallback logic and remove the duplicated code block.

final_key = _resolve_key(resolved_key, f.keys()) t = f.get_tensor(final_key)

gemini-code-assist · 2026-09-09T23:10:55Z

  if config.scan_layers:
    # If it's a standard scanned layer, we use the configured param_scan_axis.
-    axis_to_stack = config.param_scan_axis
+    axis_to_stack = config.param_scan_axis if len(target_shape) > config.param_scan_axis else 0


The function _build_single_axis_stacked_tensor is duplicated across three files: to_maxtext.py, utils/tensor_handling.py, and utils/utils.py. This duplication makes maintenance difficult and error-prone (as seen in this PR where the same fix had to be applied to all three files). Consider defining this function once in a shared utility module (e.g., utils/tensor_handling.py) and importing it where needed.

gemini-code-assist · 2026-09-09T23:10:55Z

+    elif arg.startswith("lazy_load_tensors="):
+      lazy_load_tensors = str2bool(arg.split("=", 1)[1])


Using str2bool here might raise a NameError if it is not imported or defined in this scope. A safer and more standard way to parse the boolean value from the command-line argument string is to check if it is in ('true', '1').

Suggested change

elif arg.startswith("lazy_load_tensors="):

lazy_load_tensors = str2bool(arg.split("=", 1)[1])

elif arg.startswith("lazy_load_tensors="):

lazy_load_tensors = arg.split("=", 1)[1].lower() in ("true", "1")

gemini-code-assist · 2026-09-09T23:10:55Z

+          if torch is not None and isinstance(v, torch.Tensor):
+            if v.dtype == torch.bfloat16:
+              v = v.to(torch.float32)
+            return v.to(torch.float32).cpu().numpy().astype(ml_dtypes.bfloat16)


The check if v.dtype == torch.bfloat16: v = v.to(torch.float32) is redundant because v.to(torch.float32) is called unconditionally on the very next line in the return statement. You can simplify this to a single line.

if torch is not None and isinstance(v, torch.Tensor): return v.to(torch.float32).cpu().numpy().astype(ml_dtypes.bfloat16)

gemini-code-assist · 2026-09-09T23:10:55Z

+    "load_parameters_path=/dev/shm/maxtext_qwen3.5_35b_fp8/0/items",
+    "scan_layers=true",
+    "per_device_batch_size=1",
+    "max_prefill_predict_length=4",
+    "max_target_length=4",
+    "async_checkpointing=false",
+    "sparse_matmul=false",
+    "ici_fsdp_parallelism=1",
+    "ici_expert_parallelism=-1",
+    "matmul_precision=highest",
+    "float32_logits=true",
+    "float32_qk_product=true",
+    "--golden_logits_path=/dev/shm/golden_qwen3.5_35b_fp8.pkl",
+    "--max_kl_div=0.2",


The test script hardcodes absolute paths in /dev/shm/ (e.g., /dev/shm/maxtext_qwen3.5_35b_fp8/0/items and /dev/shm/golden_qwen3.5_35b_fp8.pkl). These paths are environment-specific and will cause the script to fail on other systems or in CI/CD pipelines. Consider using environment variables or command-line arguments to specify these paths, or at least document how to set up these files.

gobbleturk · 2026-09-10T17:26:33Z

+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# model config for qwen3.5-397b-a17b-fp8 (FP8 weight-only storage with dynamic dequantization)


what does dynamic dequant mean? Aren't the scales known ahead of time?

gobbleturk

Can you include performance/xprof links in the PR description? I thought block scaling was inefficient except for blocks>512 - at least block=256 to fit our MXU shape?

snehalv2002 added 11 commits September 2, 2026 22:25

feat(layers): implement FP8 weight-only dynamic dequantization engine…

f74726d

… and types

feat(layers): support generic weight_block_size configuration for blo…

9f2d431

…ck-scaled quantization

refactor(layers): clarify DenseGeneral scale casework and restore att…

f77f5ef

…end_dtype

refactor(engine): make core dequantization layers and types model-agn…

12efee3

…ostic

refactor(common_types): isolate unquantized module resolution entirel…

27b90ec

…y to config

feat(checkpoint_conversion): support direct FP8 and scale tensor inge…

c9d9613

…stion in to_maxtext

fix(checkpoint_conversion): normalize CLI arguments and keyword overr…

f54bfa9

…ides in to_maxtext

feat(checkpointing): add dequantize-on-load parameter restoration

c427c2c

feat(models): add qwen3.5-35b-a3b-fp8 model config and HuggingFace ma…

ada785a

…pping

feat(models): wire weight_block_size config for qwen3.5-35b-a3b-fp8

99908da

Shuwen-Fang requested review from NuojCheng, RissyRan, abhinavclemson, aireenmei, bvandermoon, darisoy, gagika, gobbleturk, hengtaoguo, huytransformer, igorts-git, jiangjy1982, michelle-yooh, parambole, richjames0, shralex, shuningjin, xibinliu and zxhe-sean as code owners September 9, 2026 23:07

Shuwen-Fang requested review from A9isha, SurbhiJainUSC, dipannita08, khatwanimohit and vipannalla as code owners September 9, 2026 23:07

gemini-code-assist Bot reviewed Sep 9, 2026

View reviewed changes

Shuwen-Fang changed the base branch from main to pr/fp8-qwen3.5-35b-onboarding September 9, 2026 23:12

onboard 397b

1bb4b03

Shuwen-Fang force-pushed the qwen3.5-397b-fp8 branch from 64ad4e1 to 1bb4b03 Compare September 9, 2026 23:21

gobbleturk reviewed Sep 10, 2026

View reviewed changes

gobbleturk approved these changes Sep 10, 2026

View reviewed changes

snehalv2002 force-pushed the pr/fp8-qwen3.5-35b-onboarding branch 13 times, most recently from a1e42e1 to 93a2b66 Compare September 12, 2026 20:35

Navigation Menu

Uh oh!

Onboard Qwen3.5 397b fp8 - #5182

Onboard Qwen3.5 397b fp8#5182
Shuwen-Fang wants to merge 12 commits into
pr/fp8-qwen3.5-35b-onboardingfrom
qwen3.5-397b-fp8

Shuwen-Fang commented Sep 9, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Uh oh!

gobbleturk Sep 10, 2026

Uh oh!

gobbleturk left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		if "HF_HOME" not in os.environ and os.path.exists("/dev/shm"):
		os.environ["HF_HOME"] = "/dev/shm/hf_cache"

		elif arg.startswith("lazy_load_tensors="):
		lazy_load_tensors = str2bool(arg.split("=", 1)[1])

Uh oh!

Conversation

Shuwen-Fang commented Sep 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Tests

Checklist

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

gobbleturk Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

gobbleturk left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Shuwen-Fang commented Sep 9, 2026 •

edited

Loading