From 72176e9cb2216a1d600c9b2b09b3021bee6fa4a4 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Wed, 2 Sep 2026 09:50:04 -0700 Subject: [PATCH 1/6] Fix TEGroupedMLP quantizer checkpoint resharding Signed-off-by: Jennifer Chen --- .../opt/plugins/mcore_dist_checkpointing.py | 19 ++ .../torch/quantization/plugins/megatron.py | 7 + .../quantization/plugins/test_megatron.py | 179 ++++++++++-------- 3 files changed, 131 insertions(+), 74 deletions(-) diff --git a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py index 2bae53fa6f1..da114825c8a 100644 --- a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py +++ b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py @@ -183,6 +183,24 @@ def _load_extra_state_from_sharded_checkpoint( module.set_extra_state(extra_state_dict_no_prefix[key]) +def _initialize_grouped_weight_amax_for_restore(model: torch.nn.Module) -> None: + """Create scalar amax placeholders for TE grouped experts absent on the save rank.""" + for module in model.modules(): + if not hasattr(module, "num_gemms") or not hasattr(module, "weight_quantizer"): + continue + reference_tensor = next(module.parameters(), None) + if reference_tensor is None: + continue + for gemm_idx in range(module.num_gemms): + quantizer = module.weight_quantizer[gemm_idx] + quantizers = quantizer if isinstance(quantizer, torch.nn.Sequential) else [quantizer] + for tensor_quantizer in quantizers: + if tensor_quantizer.amax is None: + tensor_quantizer.amax = torch.empty( + 1, device=reference_tensor.device, dtype=reference_tensor.dtype + ) + + def restore_sharded_modelopt_state( model: list[torch.nn.Module], checkpoint_name: str | Path, @@ -238,3 +256,4 @@ def restore_sharded_modelopt_state( model[0] = mto.restore_from_modelopt_state(model[0], common_modelopt_state) _load_extra_state_from_sharded_checkpoint(model[0], checkpoint_name, prefix, metadata=metadata) + _initialize_grouped_weight_amax_for_restore(model[0]) diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 6b3fd85b4d9..5c532c4ebc6 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -916,6 +916,13 @@ def _setup(self): self.linear_fc1._parallel_state = self.parallel_state self.linear_fc2._parallel_state = self.parallel_state + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): + """Save per-expert quantizer state as globally named singleton shards.""" + if metadata is None: + metadata = {} + metadata["singleton_local_shards"] = True + return super().sharded_state_dict(prefix, sharded_offsets, metadata) + @QuantModuleRegistry.register({TEDotProductAttention: "TEDotProductAttention"}) class _QuantTEDotProductAttention(QuantModule): """Quantized version of TEDotProductAttention for Megatron models with KV cache quantization. diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index aa8325fd328..4ce2b382eb2 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -15,7 +15,6 @@ import copy import math -import re import sys import types from contextlib import nullcontext @@ -38,7 +37,9 @@ get_batch, get_forward, initialize_for_megatron, + load_distributed_checkpoint, run_mcore_inference, + save_distributed_checkpoint, sharded_state_dict_test_helper, ) from _test_utils.torch.misc import set_seed @@ -52,6 +53,7 @@ from megatron.core.parallel_state import ( destroy_model_parallel, get_data_parallel_group, + get_expert_model_parallel_rank, get_tensor_model_parallel_group, ) from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear @@ -62,6 +64,10 @@ import modelopt import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq +from modelopt.torch.opt.plugins.mcore_dist_checkpointing import ( + restore_sharded_modelopt_state, + save_sharded_modelopt_state, +) from modelopt.torch.quantization.algorithms import QuantRecipe, _AutoQuantizeBaseSearcher from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer from modelopt.torch.quantization.nn.modules.quant_linear import RealQuantLinear @@ -1090,100 +1096,125 @@ def test_te_grouped_vs_sequential_default_amax(dist_workers_size_1, quant_cfg): ) -def _te_grouped_expert_identity_from_sharded_state(module): - """Return {local_key: (global_expert_idx, num_global_experts)} for per-expert amax shards. +def _set_te_grouped_weight_amax(model, ep_rank, num_local_experts): + """Give every local expert a distinct amax derived from its global expert index.""" + for linear in model.modules(): + if not isinstance(linear, _QuantMegatronTEGroupedLinear): + continue + for local_expert_idx in range(linear.num_gemms): + quantizer = linear.weight_quantizer[local_expert_idx] + leaves = list(quantizer) if isinstance(quantizer, SequentialQuantizer) else [quantizer] + for leaf in leaves: + if leaf._amax is not None: + leaf._amax.fill_(1.0 + ep_rank * num_local_experts + local_expert_idx) + - The grouped linear must give each fused expert the same global identity the weights use: - the dict key keeps the local expert index (maps to the local buffer on restore) while the - ShardedTensor carries the global expert offset. Called with sharded_offsets=() so the expert - axis is the (only) prepended axis at index 0. - """ - sharded_sd = module.sharded_state_dict(prefix="", sharded_offsets=(), metadata=None) - identity = {} - for key, sh_ten in sharded_sd.items(): - if re.match(r"weight_quantizer\.\d+\..*_amax$", key): - assert sh_ten.prepend_axis_num >= 1, f"{key}: expected a prepended expert axis" - identity[key] = (int(sh_ten.global_offset[0]), int(sh_ten.global_shape[0])) - return identity +def _assert_te_grouped_weight_amax(model, expected_amax): + for linear in model.modules(): + if not isinstance(linear, _QuantMegatronTEGroupedLinear): + continue + for local_expert_idx in range(linear.num_gemms): + quantizer = linear.weight_quantizer[local_expert_idx] + leaves = list(quantizer) if isinstance(quantizer, SequentialQuantizer) else [quantizer] + for leaf in leaves: + assert leaf._amax is not None + assert torch.equal( + leaf._amax, torch.full_like(leaf._amax, expected_amax[local_expert_idx]) + ) -def _test_te_grouped_sharded_state_dict_global_expert_identity_helper( - tp_size, ep_size, quant_cfg, rank, size +def _test_te_grouped_sharded_state_dict_reshard_helper( + save_tp_size, + save_ep_size, + load_tp_size, + load_ep_size, + checkpoint_path, + rank, + size, ): - """Per-expert quantizer amax must persist all num_global_experts across EP. - - With EP>1 the base linear emitted ``weight_quantizer.{local_i}._amax`` at the local index with - no expert offset, so every rank wrote identical keys and torch_dist dedup collapsed them to a - single rank's experts. Assert each rank's fused experts now carry distinct global identities so - the union across ranks covers every global expert. - """ + """Round-trip TEGroupedMLP amax through a topology change.""" + num_experts = 4 + save_num_local_experts = num_experts // save_ep_size initialize_for_megatron( - tensor_model_parallel_size=tp_size, - expert_model_parallel_size=ep_size, + tensor_model_parallel_size=save_tp_size, + expert_model_parallel_size=save_ep_size, seed=SEED, ) - num_experts = 4 - num_local = num_experts // ep_size - te_grouped = _gpt_model_provider( - tp_size=tp_size, - ep_size=ep_size, + source = _gpt_model_provider( + tp_size=save_tp_size, + ep_size=save_ep_size, hidden_size=32, moe_grouped_gemm=True, transformer_impl="transformer_engine", num_moe_experts=num_experts, ) - forward = get_forward(te_grouped, batch_size=8) - for module in te_grouped.modules(): + forward = get_forward(source, batch_size=8) + for module in source.modules(): if isinstance(module, TopKRouter): module.topk = module.num_experts - mtq.quantize(te_grouped, quant_cfg, forward) + mtq.quantize(source, mtq.NVFP4_DEFAULT_CFG, forward) + _set_te_grouped_weight_amax(source, get_expert_model_parallel_rank(), save_num_local_experts) + save_distributed_checkpoint(checkpoint_path, source) + save_sharded_modelopt_state([source], checkpoint_path) + torch.distributed.barrier() + del source + destroy_model_parallel() - grouped_linears = [ - m for m in te_grouped.modules() if isinstance(m, _QuantMegatronTEGroupedLinear) - ] - assert grouped_linears, "No grouped quant linears found" - - expected_global = {rank * num_local + i for i in range(num_local)} - for linear in grouped_linears: - # Give each expert a distinct amax so a value mix-up would also be observable. - for i in range(linear.num_gemms): - wq = linear.weight_quantizer[i] - leaves = list(wq) if isinstance(wq, SequentialQuantizer) else [wq] - for leaf in leaves: - if hasattr(leaf, "_amax") and leaf._amax is not None: - leaf._amax.fill_(1.0 + rank * num_local + i) - - identity = _te_grouped_expert_identity_from_sharded_state(linear) - # One entry per local expert per amax buffer; dict keys keep the LOCAL index. - local_keys = {int(re.search(r"weight_quantizer\.(\d+)\.", k).group(1)) for k in identity} - assert local_keys == set(range(num_local)), ( - f"Expected local expert keys {set(range(num_local))}, got {local_keys}" - ) - # ShardedTensor global identity: this rank owns experts {rank*num_local + i}. - local_global = {gidx for gidx, _ in identity.values()} - assert local_global == expected_global, ( - f"rank {rank}: expected global experts {expected_global}, got {local_global}" - ) - assert all(total == num_experts for _, total in identity.values()), ( - f"num_global_experts should be {num_experts}, got {identity}" + initialize_for_megatron( + tensor_model_parallel_size=load_tp_size, + expert_model_parallel_size=load_ep_size, + seed=SEED, + ) + target = _gpt_model_provider( + tp_size=load_tp_size, + ep_size=load_ep_size, + hidden_size=32, + moe_grouped_gemm=True, + transformer_impl="transformer_engine", + num_moe_experts=num_experts, + ) + target_models = [target] + restore_sharded_modelopt_state(target_models, checkpoint_path) + target = target_models[0] + load_distributed_checkpoint(checkpoint_path, target) + load_num_local_experts = num_experts // load_ep_size + expected_amax = tuple( + range( + 1 + get_expert_model_parallel_rank() * load_num_local_experts, + 1 + (get_expert_model_parallel_rank() + 1) * load_num_local_experts, ) - - # Gather the global expert indices across all EP ranks: the union must cover every expert. - gathered = [None] * size - torch.distributed.all_gather_object(gathered, sorted(expected_global)) - union = set() - for part in gathered: - union.update(part) - assert union == set(range(num_experts)), ( - f"Union of global experts across EP ranks should be {set(range(num_experts))}, got {union}" ) + _assert_te_grouped_weight_amax(target, expected_amax) -@pytest.mark.parametrize("quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG]) -def test_te_grouped_sharded_state_dict_global_expert_identity(dist_workers_size_2, quant_cfg): +@pytest.mark.parametrize( + ( + "save_tp_size", + "save_ep_size", + "load_tp_size", + "load_ep_size", + ), + [(1, 2, 1, 1), (1, 1, 1, 2), (1, 1, 2, 1), (2, 1, 1, 1)], + ids=["ep-downsize", "ep-upsize", "tp-upsize", "tp-downsize"], +) +def test_te_grouped_sharded_state_dict_reshard( + dist_workers_size_2, + tmp_path, + save_tp_size, + save_ep_size, + load_tp_size, + load_ep_size, +): dist_workers_size_2.run( - partial(_test_te_grouped_sharded_state_dict_global_expert_identity_helper, 1, 2, quant_cfg) + partial( + _test_te_grouped_sharded_state_dict_reshard_helper, + save_tp_size, + save_ep_size, + load_tp_size, + load_ep_size, + tmp_path, + ) ) From 5d2ebeac47d2879692516744ddf03843a5bb4e86 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Thu, 3 Sep 2026 07:06:40 -0700 Subject: [PATCH 2/6] Use deterministic TE grouped amax placeholders Signed-off-by: Jennifer Chen --- .../opt/plugins/mcore_dist_checkpointing.py | 22 ++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py index da114825c8a..d9cb9840278 100644 --- a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py +++ b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py @@ -184,7 +184,7 @@ def _load_extra_state_from_sharded_checkpoint( def _initialize_grouped_weight_amax_for_restore(model: torch.nn.Module) -> None: - """Create scalar amax placeholders for TE grouped experts absent on the save rank.""" + """Create deterministic amax placeholders for TE grouped experts absent on the save rank.""" for module in model.modules(): if not hasattr(module, "num_gemms") or not hasattr(module, "weight_quantizer"): continue @@ -194,11 +194,23 @@ def _initialize_grouped_weight_amax_for_restore(model: torch.nn.Module) -> None: for gemm_idx in range(module.num_gemms): quantizer = module.weight_quantizer[gemm_idx] quantizers = quantizer if isinstance(quantizer, torch.nn.Sequential) else [quantizer] - for tensor_quantizer in quantizers: + for quantizer_idx, tensor_quantizer in enumerate(quantizers): if tensor_quantizer.amax is None: - tensor_quantizer.amax = torch.empty( - 1, device=reference_tensor.device, dtype=reference_tensor.dtype - ) + for sibling_gemm_idx in range(module.num_gemms): + sibling_quantizer = module.weight_quantizer[sibling_gemm_idx] + sibling_quantizers = ( + sibling_quantizer + if isinstance(sibling_quantizer, torch.nn.Sequential) + else [sibling_quantizer] + ) + if quantizer_idx >= len(sibling_quantizers): + continue + sibling_amax = sibling_quantizers[quantizer_idx].amax + if sibling_amax is not None: + tensor_quantizer.amax = torch.zeros_like( + sibling_amax, device=reference_tensor.device + ) + break def restore_sharded_modelopt_state( From 086bd21a2504e1a1db4d63de6e9e2526e4fb952b Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Thu, 3 Sep 2026 10:05:00 -0700 Subject: [PATCH 3/6] Preserve TE grouped quantizer restore state Signed-off-by: Jennifer Chen --- .../opt/plugins/mcore_dist_checkpointing.py | 28 ++++++++++++----- .../quantization/plugins/test_megatron.py | 30 +++++++++++++++++++ 2 files changed, 50 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py index d9cb9840278..160fc92b71d 100644 --- a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py +++ b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py @@ -183,8 +183,8 @@ def _load_extra_state_from_sharded_checkpoint( module.set_extra_state(extra_state_dict_no_prefix[key]) -def _initialize_grouped_weight_amax_for_restore(model: torch.nn.Module) -> None: - """Create deterministic amax placeholders for TE grouped experts absent on the save rank.""" +def _initialize_grouped_weight_quantizer_state_for_restore(model: torch.nn.Module) -> None: + """Create deterministic quantizer-state placeholders for TE grouped experts.""" for module in model.modules(): if not hasattr(module, "num_gemms") or not hasattr(module, "weight_quantizer"): continue @@ -195,7 +195,14 @@ def _initialize_grouped_weight_amax_for_restore(model: torch.nn.Module) -> None: quantizer = module.weight_quantizer[gemm_idx] quantizers = quantizer if isinstance(quantizer, torch.nn.Sequential) else [quantizer] for quantizer_idx, tensor_quantizer in enumerate(quantizers): - if tensor_quantizer.amax is None: + if not tensor_quantizer.is_enabled or tensor_quantizer.is_mx_format: + continue + for state_name in ("amax", "global_amax"): + if ( + not hasattr(tensor_quantizer, state_name) + or getattr(tensor_quantizer, state_name) is not None + ): + continue for sibling_gemm_idx in range(module.num_gemms): sibling_quantizer = module.weight_quantizer[sibling_gemm_idx] sibling_quantizers = ( @@ -205,10 +212,15 @@ def _initialize_grouped_weight_amax_for_restore(model: torch.nn.Module) -> None: ) if quantizer_idx >= len(sibling_quantizers): continue - sibling_amax = sibling_quantizers[quantizer_idx].amax - if sibling_amax is not None: - tensor_quantizer.amax = torch.zeros_like( - sibling_amax, device=reference_tensor.device + sibling_quantizer = sibling_quantizers[quantizer_idx] + if not sibling_quantizer.is_enabled or sibling_quantizer.is_mx_format: + continue + sibling_state = getattr(sibling_quantizer, state_name, None) + if sibling_state is not None: + setattr( + tensor_quantizer, + state_name, + torch.zeros_like(sibling_state, device=reference_tensor.device), ) break @@ -268,4 +280,4 @@ def restore_sharded_modelopt_state( model[0] = mto.restore_from_modelopt_state(model[0], common_modelopt_state) _load_extra_state_from_sharded_checkpoint(model[0], checkpoint_name, prefix, metadata=metadata) - _initialize_grouped_weight_amax_for_restore(model[0]) + _initialize_grouped_weight_quantizer_state_for_restore(model[0]) diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index 4ce2b382eb2..2aeb9209fd9 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -65,6 +65,7 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.opt.plugins.mcore_dist_checkpointing import ( + _initialize_grouped_weight_quantizer_state_for_restore, restore_sharded_modelopt_state, save_sharded_modelopt_state, ) @@ -1123,6 +1124,35 @@ def _assert_te_grouped_weight_amax(model, expected_amax): ) +def test_initialize_grouped_weight_quantizer_state_for_restore(): + """Missing grouped state inherits the shape and dtype of a populated sibling.""" + source = mtq.nn.StaticBlockScaleQuantizer.from_tensor_quantizer( + mtq.nn.TensorQuantizer(amax=torch.tensor([1.0, 2.0])), global_amax=torch.tensor(2.0) + ) + target = mtq.nn.StaticBlockScaleQuantizer.from_tensor_quantizer(mtq.nn.TensorQuantizer()) + disabled = mtq.nn.StaticBlockScaleQuantizer.from_tensor_quantizer(mtq.nn.TensorQuantizer()) + disabled.disable() + mx = mtq.nn.TensorQuantizer( + mtq.config.QuantizerAttributeConfig( + num_bits=(4, 3), block_sizes={-1: 32, "type": "dynamic", "scale_bits": (8, 0)} + ) + ) + + model = torch.nn.Module() + model.weight = torch.nn.Parameter(torch.empty(1)) + model.num_gemms = 4 + model.weight_quantizer = torch.nn.ModuleList([source, target, disabled, mx]) + + _initialize_grouped_weight_quantizer_state_for_restore(model) + + assert torch.equal(target.amax, torch.zeros_like(source.amax)) + assert torch.equal(target.global_amax, torch.zeros_like(source.global_amax)) + assert disabled.amax is None + assert disabled.global_amax is None + assert mx.amax is None + assert not hasattr(mx, "_amax") + + def _test_te_grouped_sharded_state_dict_reshard_helper( save_tp_size, save_ep_size, From a7947509710ca33f41f825768768f9148366da10 Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Thu, 3 Sep 2026 10:30:38 -0700 Subject: [PATCH 4/6] Fix TE grouped quantizer resharding Signed-off-by: Jennifer Chen --- .../opt/plugins/mcore_dist_checkpointing.py | 52 +++++------- .../torch/quantization/plugins/megatron.py | 7 -- .../quantization/plugins/test_megatron.py | 79 ++++++++++++++++--- 3 files changed, 84 insertions(+), 54 deletions(-) diff --git a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py index 160fc92b71d..2de993a5371 100644 --- a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py +++ b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py @@ -188,41 +188,25 @@ def _initialize_grouped_weight_quantizer_state_for_restore(model: torch.nn.Modul for module in model.modules(): if not hasattr(module, "num_gemms") or not hasattr(module, "weight_quantizer"): continue - reference_tensor = next(module.parameters(), None) - if reference_tensor is None: - continue - for gemm_idx in range(module.num_gemms): - quantizer = module.weight_quantizer[gemm_idx] - quantizers = quantizer if isinstance(quantizer, torch.nn.Sequential) else [quantizer] - for quantizer_idx, tensor_quantizer in enumerate(quantizers): - if not tensor_quantizer.is_enabled or tensor_quantizer.is_mx_format: + grouped_leaves = [ + quantizer if isinstance(quantizer, torch.nn.Sequential) else [quantizer] + for quantizer in [module.weight_quantizer[idx] for idx in range(module.num_gemms)] + ] + for sibling_leaves in zip(*grouped_leaves): + eligible_leaves = [ + quantizer + for quantizer in sibling_leaves + if quantizer.is_enabled and not quantizer.is_mx_format + ] + for state_name in ("amax", "global_amax"): + reference = next( + (getattr(quantizer, state_name, None) for quantizer in eligible_leaves), None + ) + if reference is None: continue - for state_name in ("amax", "global_amax"): - if ( - not hasattr(tensor_quantizer, state_name) - or getattr(tensor_quantizer, state_name) is not None - ): - continue - for sibling_gemm_idx in range(module.num_gemms): - sibling_quantizer = module.weight_quantizer[sibling_gemm_idx] - sibling_quantizers = ( - sibling_quantizer - if isinstance(sibling_quantizer, torch.nn.Sequential) - else [sibling_quantizer] - ) - if quantizer_idx >= len(sibling_quantizers): - continue - sibling_quantizer = sibling_quantizers[quantizer_idx] - if not sibling_quantizer.is_enabled or sibling_quantizer.is_mx_format: - continue - sibling_state = getattr(sibling_quantizer, state_name, None) - if sibling_state is not None: - setattr( - tensor_quantizer, - state_name, - torch.zeros_like(sibling_state, device=reference_tensor.device), - ) - break + for quantizer in eligible_leaves: + if getattr(quantizer, state_name, None) is None: + setattr(quantizer, state_name, torch.zeros_like(reference)) def restore_sharded_modelopt_state( diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 5c532c4ebc6..6b3fd85b4d9 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -916,13 +916,6 @@ def _setup(self): self.linear_fc1._parallel_state = self.parallel_state self.linear_fc2._parallel_state = self.parallel_state - def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): - """Save per-expert quantizer state as globally named singleton shards.""" - if metadata is None: - metadata = {} - metadata["singleton_local_shards"] = True - return super().sharded_state_dict(prefix, sharded_offsets, metadata) - @QuantModuleRegistry.register({TEDotProductAttention: "TEDotProductAttention"}) class _QuantTEDotProductAttention(QuantModule): """Quantized version of TEDotProductAttention for Megatron models with KV cache quantization. diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index 2aeb9209fd9..a00f6b22026 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -1097,8 +1097,8 @@ def test_te_grouped_vs_sequential_default_amax(dist_workers_size_1, quant_cfg): ) -def _set_te_grouped_weight_amax(model, ep_rank, num_local_experts): - """Give every local expert a distinct amax derived from its global expert index.""" +def _set_te_grouped_weight_quantizer_state(model, ep_rank, num_local_experts): + """Give every local expert distinct quantizer state derived from its global index.""" for linear in model.modules(): if not isinstance(linear, _QuantMegatronTEGroupedLinear): continue @@ -1106,11 +1106,16 @@ def _set_te_grouped_weight_amax(model, ep_rank, num_local_experts): quantizer = linear.weight_quantizer[local_expert_idx] leaves = list(quantizer) if isinstance(quantizer, SequentialQuantizer) else [quantizer] for leaf in leaves: - if leaf._amax is not None: - leaf._amax.fill_(1.0 + ep_rank * num_local_experts + local_expert_idx) + amax = getattr(leaf, "_amax", None) + if amax is not None: + amax.fill_(1.0 + ep_rank * num_local_experts + local_expert_idx) + global_amax = getattr(leaf, "_global_amax", None) + if global_amax is not None: + global_amax.fill_(1.0 + ep_rank * num_local_experts + local_expert_idx) -def _assert_te_grouped_weight_amax(model, expected_amax): +def _assert_te_grouped_weight_quantizer_state(model, expected_amax, expect_global_amax): + checked = 0 for linear in model.modules(): if not isinstance(linear, _QuantMegatronTEGroupedLinear): continue @@ -1118,10 +1123,22 @@ def _assert_te_grouped_weight_amax(model, expected_amax): quantizer = linear.weight_quantizer[local_expert_idx] leaves = list(quantizer) if isinstance(quantizer, SequentialQuantizer) else [quantizer] for leaf in leaves: - assert leaf._amax is not None - assert torch.equal( - leaf._amax, torch.full_like(leaf._amax, expected_amax[local_expert_idx]) + amax = getattr(leaf, "_amax", None) + assert amax is not None, ( + "TEGrouped per-expert weight quantizer amax was not restored" ) + checked += 1 + assert torch.equal(amax, torch.full_like(amax, expected_amax[local_expert_idx])) + global_amax = getattr(leaf, "_global_amax", None) + if expect_global_amax: + assert global_amax is not None, ( + "TEGrouped per-expert weight quantizer global_amax was not restored" + ) + assert torch.equal( + global_amax, + torch.full_like(global_amax, expected_amax[local_expert_idx]), + ) + assert checked > 0, "no TEGrouped per-expert weight quantizer amax was checked" def test_initialize_grouped_weight_quantizer_state_for_restore(): @@ -1158,6 +1175,8 @@ def _test_te_grouped_sharded_state_dict_reshard_helper( save_ep_size, load_tp_size, load_ep_size, + quant_cfg, + expect_global_amax, checkpoint_path, rank, size, @@ -1183,8 +1202,10 @@ def _test_te_grouped_sharded_state_dict_reshard_helper( for module in source.modules(): if isinstance(module, TopKRouter): module.topk = module.num_experts - mtq.quantize(source, mtq.NVFP4_DEFAULT_CFG, forward) - _set_te_grouped_weight_amax(source, get_expert_model_parallel_rank(), save_num_local_experts) + mtq.quantize(source, copy.deepcopy(quant_cfg), forward) + _set_te_grouped_weight_quantizer_state( + source, get_expert_model_parallel_rank(), save_num_local_experts + ) save_distributed_checkpoint(checkpoint_path, source) save_sharded_modelopt_state([source], checkpoint_path) torch.distributed.barrier() @@ -1215,18 +1236,46 @@ def _test_te_grouped_sharded_state_dict_reshard_helper( 1 + (get_expert_model_parallel_rank() + 1) * load_num_local_experts, ) ) - _assert_te_grouped_weight_amax(target, expected_amax) + _assert_te_grouped_weight_quantizer_state(target, expected_amax, expect_global_amax) @pytest.mark.parametrize( ( + "quant_cfg", + "expect_global_amax", "save_tp_size", "save_ep_size", "load_tp_size", "load_ep_size", ), - [(1, 2, 1, 1), (1, 1, 1, 2), (1, 1, 2, 1), (2, 1, 1, 1)], - ids=["ep-downsize", "ep-upsize", "tp-upsize", "tp-downsize"], + [ + pytest.param(mtq.FP8_DEFAULT_CFG, False, 1, 2, 1, 1, id="fp8-ep-downsize"), + pytest.param(mtq.FP8_DEFAULT_CFG, False, 1, 1, 1, 2, id="fp8-ep-upsize"), + pytest.param(mtq.FP8_DEFAULT_CFG, False, 1, 1, 2, 1, id="fp8-tp-upsize"), + pytest.param(mtq.FP8_DEFAULT_CFG, False, 2, 1, 1, 1, id="fp8-tp-downsize"), + pytest.param(mtq.NVFP4_DEFAULT_CFG, False, 1, 2, 1, 1, id="nvfp4-ep-downsize"), + pytest.param(mtq.NVFP4_DEFAULT_CFG, False, 1, 1, 1, 2, id="nvfp4-ep-upsize"), + pytest.param(mtq.NVFP4_DEFAULT_CFG, False, 1, 1, 2, 1, id="nvfp4-tp-upsize"), + pytest.param(mtq.NVFP4_DEFAULT_CFG, False, 2, 1, 1, 1, id="nvfp4-tp-downsize"), + pytest.param( + mtq.NVFP4_W4A4_WEIGHT_MSE_FP8_SWEEP_CFG, + True, + 1, + 2, + 1, + 1, + id="nvfp4-mse-ep-downsize", + ), + pytest.param( + mtq.NVFP4_W4A4_WEIGHT_MSE_FP8_SWEEP_CFG, + True, + 1, + 1, + 1, + 2, + id="nvfp4-mse-ep-upsize", + ), + ], ) def test_te_grouped_sharded_state_dict_reshard( dist_workers_size_2, @@ -1235,6 +1284,8 @@ def test_te_grouped_sharded_state_dict_reshard( save_ep_size, load_tp_size, load_ep_size, + quant_cfg, + expect_global_amax, ): dist_workers_size_2.run( partial( @@ -1243,6 +1294,8 @@ def test_te_grouped_sharded_state_dict_reshard( save_ep_size, load_tp_size, load_ep_size, + quant_cfg, + expect_global_amax, tmp_path, ) ) From 3ceb427bbd69131008a36bcdbe9887189e135bbb Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Thu, 3 Sep 2026 12:26:57 -0700 Subject: [PATCH 5/6] Handle missing grouped quantizer siblings Signed-off-by: Jennifer Chen --- modelopt/torch/opt/plugins/mcore_dist_checkpointing.py | 7 ++++++- .../torch/quantization/plugins/test_megatron.py | 2 +- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py index 2de993a5371..e2cbe028e01 100644 --- a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py +++ b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py @@ -200,7 +200,12 @@ def _initialize_grouped_weight_quantizer_state_for_restore(model: torch.nn.Modul ] for state_name in ("amax", "global_amax"): reference = next( - (getattr(quantizer, state_name, None) for quantizer in eligible_leaves), None + ( + state + for quantizer in eligible_leaves + if (state := getattr(quantizer, state_name, None)) is not None + ), + None, ) if reference is None: continue diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index a00f6b22026..b81b2572837 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -1158,7 +1158,7 @@ def test_initialize_grouped_weight_quantizer_state_for_restore(): model = torch.nn.Module() model.weight = torch.nn.Parameter(torch.empty(1)) model.num_gemms = 4 - model.weight_quantizer = torch.nn.ModuleList([source, target, disabled, mx]) + model.weight_quantizer = torch.nn.ModuleList([target, source, disabled, mx]) _initialize_grouped_weight_quantizer_state_for_restore(model) From ab7fd65fec55796790806552e152620dba9c971b Mon Sep 17 00:00:00 2001 From: Jennifer Chen Date: Fri, 4 Sep 2026 08:24:00 -0700 Subject: [PATCH 6/6] Move grouped restore state into quantization plugin Signed-off-by: Jennifer Chen --- .../opt/plugins/mcore_dist_checkpointing.py | 34 ++---------------- .../torch/quantization/plugins/megatron.py | 36 +++++++++++++++++++ .../quantization/plugins/test_megatron.py | 4 +-- 3 files changed, 41 insertions(+), 33 deletions(-) diff --git a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py index e2cbe028e01..a4796162f6d 100644 --- a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py +++ b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py @@ -181,37 +181,10 @@ def _load_extra_state_from_sharded_checkpoint( module, "modelopt_set_extra_state_callbacks" ): module.set_extra_state(extra_state_dict_no_prefix[key]) - - -def _initialize_grouped_weight_quantizer_state_for_restore(model: torch.nn.Module) -> None: - """Create deterministic quantizer-state placeholders for TE grouped experts.""" for module in model.modules(): - if not hasattr(module, "num_gemms") or not hasattr(module, "weight_quantizer"): - continue - grouped_leaves = [ - quantizer if isinstance(quantizer, torch.nn.Sequential) else [quantizer] - for quantizer in [module.weight_quantizer[idx] for idx in range(module.num_gemms)] - ] - for sibling_leaves in zip(*grouped_leaves): - eligible_leaves = [ - quantizer - for quantizer in sibling_leaves - if quantizer.is_enabled and not quantizer.is_mx_format - ] - for state_name in ("amax", "global_amax"): - reference = next( - ( - state - for quantizer in eligible_leaves - if (state := getattr(quantizer, state_name, None)) is not None - ), - None, - ) - if reference is None: - continue - for quantizer in eligible_leaves: - if getattr(quantizer, state_name, None) is None: - setattr(quantizer, state_name, torch.zeros_like(reference)) + post_load_extra_state = getattr(module, "modelopt_post_load_extra_state", None) + if callable(post_load_extra_state): + post_load_extra_state() def restore_sharded_modelopt_state( @@ -269,4 +242,3 @@ def restore_sharded_modelopt_state( model[0] = mto.restore_from_modelopt_state(model[0], common_modelopt_state) _load_extra_state_from_sharded_checkpoint(model[0], checkpoint_name, prefix, metadata=metadata) - _initialize_grouped_weight_quantizer_state_for_restore(model[0]) diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py index 6b3fd85b4d9..dc6c4ae2ffe 100644 --- a/modelopt/torch/quantization/plugins/megatron.py +++ b/modelopt/torch/quantization/plugins/megatron.py @@ -110,6 +110,39 @@ def _check_nvfp4_static_tp_supported(model: torch.nn.Module) -> None: ) +def _initialize_grouped_weight_quantizer_state(module: torch.nn.Module) -> None: + """Create per-expert buffers as destinations for the subsequent checkpoint load.""" + grouped_leaves = [ + quantizer if isinstance(quantizer, torch.nn.Sequential) else [quantizer] + for quantizer in [module.weight_quantizer[idx] for idx in range(module.num_gemms)] + ] + for sibling_leaves in zip(*grouped_leaves): + eligible_leaves = [ + quantizer + for quantizer in sibling_leaves + if ( + quantizer.is_enabled + and not quantizer.is_mx_format + and not getattr(quantizer, "_dynamic", False) + and not getattr(quantizer, "_lsq", False) + ) + ] + for state_name in ("_amax", "_global_amax"): + reference = next( + ( + state + for quantizer in eligible_leaves + if (state := getattr(quantizer, state_name, None)) is not None + ), + None, + ) + if reference is None: + continue + for quantizer in eligible_leaves: + if getattr(quantizer, state_name, None) is None: + quantizer.register_buffer(state_name, torch.zeros_like(reference)) + + def real_quant_module_get_extra_state(self) -> dict: """Populating real_quantizer_state and q_tensor_state.""" extra_state = {} @@ -757,6 +790,9 @@ class _QuantTELayerNormColumnParallelLinear( # Quantized subclasses to support TEGroupedLinear quantization class _QuantMegatronTEGroupedLinear(_QuantTEGroupedLinear, _MegatronParallelLinear): + def modelopt_post_load_extra_state(self): + _initialize_grouped_weight_quantizer_state(self) + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): # _sharded_state_dict_grouped adds _extra_state{gemm_idx} for gemm_idx:[1, num_gemms] in # sharded_state_dict which is same as _extra_state. The _extra_state{gemm_idx} is used for diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index b81b2572837..dc45830e816 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -65,7 +65,6 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.opt.plugins.mcore_dist_checkpointing import ( - _initialize_grouped_weight_quantizer_state_for_restore, restore_sharded_modelopt_state, save_sharded_modelopt_state, ) @@ -73,6 +72,7 @@ from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer from modelopt.torch.quantization.nn.modules.quant_linear import RealQuantLinear from modelopt.torch.quantization.plugins.megatron import ( + _initialize_grouped_weight_quantizer_state, _output_layer_untied, _QuantMegatronTEGroupedLinear, _QuantTEMCoreRowParallelLinear, @@ -1160,7 +1160,7 @@ def test_initialize_grouped_weight_quantizer_state_for_restore(): model.num_gemms = 4 model.weight_quantizer = torch.nn.ModuleList([target, source, disabled, mx]) - _initialize_grouped_weight_quantizer_state_for_restore(model) + _initialize_grouped_weight_quantizer_state(model) assert torch.equal(target.amax, torch.zeros_like(source.amax)) assert torch.equal(target.global_amax, torch.zeros_like(source.global_amax))