From e5d84e83bea9e77ee9033a7d0e9e7adaa7e021ed Mon Sep 17 00:00:00 2001 From: Boxiang Wang Date: Tue, 11 Nov 2025 20:22:54 -0800 Subject: [PATCH 1/4] Init change Signed-off-by: Boxiang Wang --- megatron/core/dist_checkpointing/dict_utils.py | 11 +++++++---- megatron/core/dist_checkpointing/serialization.py | 3 ++- megatron/core/optimizer/muon.py | 2 ++ megatron/core/optimizer/optimizer.py | 10 +++++++++- megatron/training/checkpointing.py | 2 +- 5 files changed, 21 insertions(+), 7 deletions(-) diff --git a/megatron/core/dist_checkpointing/dict_utils.py b/megatron/core/dist_checkpointing/dict_utils.py index eecf7674ef0..c40de47d734 100644 --- a/megatron/core/dist_checkpointing/dict_utils.py +++ b/megatron/core/dist_checkpointing/dict_utils.py @@ -217,22 +217,25 @@ def dict_list_map_outplace(f: Callable[[U], V], x: Union[Dict, List, U]) -> Unio return f(x) -def merge(x1: Union[dict, list], x2: Union[dict, list], key: Tuple[Union[str, int], ...] = ()): +def merge(x1: Union[dict, list], x2: Union[dict, list], key: Tuple[Union[str, int], ...] = (), layerwise_dist_opt: bool = False): """Merges dicts and lists recursively.""" if isinstance(x1, dict) and isinstance(x2, dict): for k, v2 in x2.items(): if k not in x1: x1[k] = v2 else: - x1[k] = merge(x1[k], v2, key=key + (k,)) + x1[k] = merge(x1[k], v2, key=key + (k,), layerwise_dist_opt=layerwise_dist_opt) elif isinstance(x1, list) and isinstance(x2, list): - if len(x1) != len(x2): + # for layerwise dist opt, if the common list is empty, return the loaded list + if layerwise_dist_opt and len(x1) == 0: + return x2 + elif len(x1) != len(x2): raise ValueError( f"Cannot merge two lists with different lengths ({len(x1)} and {len(x2)}, " f"encountered at level {key})" ) for i, v2 in enumerate(x2): - x1[i] = merge(x1[i], v2, key=key + (i,)) + x1[i] = merge(x1[i], v2, key=key + (i,), layerwise_dist_opt=layerwise_dist_opt) else: raise ValueError( f"Duplicate non-dict and non-list values encountered: `{x1}` and `{x2}` " diff --git a/megatron/core/dist_checkpointing/serialization.py b/megatron/core/dist_checkpointing/serialization.py index 0469949c67d..321a0d406de 100644 --- a/megatron/core/dist_checkpointing/serialization.py +++ b/megatron/core/dist_checkpointing/serialization.py @@ -64,6 +64,7 @@ def load( common_strategy: Union[LoadCommonStrategy, Tuple[str, int], None] = None, validate_access_integrity: bool = True, strict: Union[str, StrictHandling] = StrictHandling.ASSUME_OK_UNEXPECTED, + layerwise_dist_opt: bool = False, ) -> Union[StateDict, Tuple[StateDict, Set[str], Set[str]]]: """Loading entrypoint. @@ -160,7 +161,7 @@ def load( loaded_state_dict = sharded_strategy.load(sharded_state_dict, checkpoint_dir) - merge(common_state_dict, loaded_state_dict) + merge(common_state_dict, loaded_state_dict, layerwise_dist_opt=True) loaded_state_dict = apply_factory_merges(common_state_dict, sh_ten_factories) diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py index a31c84a6e8a..9fe132ea1d3 100644 --- a/megatron/core/optimizer/muon.py +++ b/megatron/core/optimizer/muon.py @@ -50,6 +50,7 @@ def __init__( use_nesterov: bool = True, weight_decay: float = 0.01, use_decoupled_weight_decay: bool = True, + use_independent_wd: bool = False, split_qkv: bool = False, is_qkv_fn: Callable[[torch.Tensor], bool] | None = None, qkv_split_shapes: tuple[int, int, int] | None = None, @@ -102,6 +103,7 @@ def scaled_orthogonalize_fn( use_nesterov, weight_decay, use_decoupled_weight_decay, + use_independent_wd, fp32_matmul_prec, scaled_orthogonalize_fn, ) diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index 1829cb424f1..f3e256f6da7 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -1222,7 +1222,15 @@ def load_state_dict(self, state_dict): if isinstance(state_dict, dict): state_dict = (v for k, v in sorted(state_dict.items())) for optimizer, state in zip(self.chained_optimizers, state_dict): - optimizer.load_state_dict(state) + # Mostly for Layerwise Dist Opt situation but it can be a general check + # if the state dict optimizer is in this rank, load the state dict + maybe_optimizer_state = state['optimizer'] if 'optimizer' in state else state['optimizer_state_dict'] + if 'state' in maybe_optimizer_state: + pass + # optimizer.load_state_dict(state) + # else: + # print(f"state_dict: {torch.distributed.get_rank()} {state['optimizer']} {type(optimizer)} \n") + # assert False self._synchronize_steps() @torch.no_grad() diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index 6ddf9f9196d..772bf1a54fd 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -1021,7 +1021,7 @@ def _load_global_dist_base_checkpoint( ) if checkpointing_context is not None: checkpointing_context["load_strategy"] = load_strategy - state_dict = dist_checkpointing.load(sharded_state_dict, checkpoint_name, load_strategy, strict=args.dist_ckpt_strictness) + state_dict = dist_checkpointing.load(sharded_state_dict, checkpoint_name, load_strategy, strict=args.dist_ckpt_strictness, layerwise_dist_opt=args.optimizer == 'dist_muon') return state_dict, checkpoint_name, release, CheckpointType.GLOBAL From 41938e9f3356619a1b7e94076972d256f9480d08 Mon Sep 17 00:00:00 2001 From: Boxiang Wang Date: Wed, 12 Nov 2025 10:50:31 -0800 Subject: [PATCH 2/4] Run experiments Signed-off-by: Boxiang Wang --- .../core/extensions/transformer_engine.py | 24 +++++++++++++++++++ megatron/training/arguments.py | 13 ++++++++++ megatron/training/checkpointing.py | 14 ++++++++++- 3 files changed, 50 insertions(+), 1 deletion(-) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index e807ee54fbf..33aa4b76b58 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -990,6 +990,14 @@ def __init__( self.kept_packed_seq_params.discard("cu_seqlens_q_padded") self.kept_packed_seq_params.discard("cu_seqlens_kv_padded") + if config.qk_clip or config.log_max_attention_logit: + # qk-clip is only supported in TE 2.9.0 and later + # assert is_te_min_version("2.9.0"), "qk-clip is only supported in TE 2.9.0 and later" + + # TE 2.9.0 introduces return_max_logit for qk-clip getting the max attention logits + extra_kwargs["return_max_logit"] = True + self.current_max_attn_logits = None + super().__init__( num_attention_heads=self.config.num_attention_heads, kv_channels=kv_channels, @@ -1059,6 +1067,22 @@ def forward( **attention_bias_kwargs, **packed_seq_kwargs, ) + + if self.config.qk_clip or self.config.log_max_attention_logit: + # qk-clip is only supported in TE 2.9.0 and later + # assert is_te_min_version("2.9.0"), "qk-clip is only supported in TE 2.9.0 and later" + + # Update Q K outside of TE Attention API + core_attn_out, batch_max_attention_logits = core_attn_out + + # Update QK_Clip balancing eta + if self.current_max_attn_logits is None: + self.current_max_attn_logits = batch_max_attention_logits + else: + self.current_max_attn_logits = torch.max( + self.current_max_attn_logits, batch_max_attention_logits + ) + else: core_attn_out = super().forward( query, key, value, attention_mask, **attention_bias_kwargs, **packed_seq_kwargs diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 3056f2007f2..4c7847f3927 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1003,6 +1003,19 @@ def validate_args(args, defaults={}): if args.add_bias_linear: args.add_qkv_bias = True + if args.qk_clip: + # assert is_te_min_version("2.9.0"), \ + # '--qk-clip is only supported with TE >= 2.9.0.' + assert 0.0 < args.qk_clip_alpha < 1.0, \ + '--qk-clip-alpha must be between 0.0 and 1.0 when using --qk-clip.' + assert args.qk_clip_threshold > 0, \ + '--qk-clip-threshold must be greater than 0 when using --qk-clip.' + + # decoupled log max attention logit check + # if args.log_max_attention_logit: + # assert is_te_min_version("2.9.0"), \ + # '--log-max-attention-logit is only supported with TE >= 2.9.0.' + # Retro checks. if args.retro_add_retriever: diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index 772bf1a54fd..24253f194bf 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -478,6 +478,14 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati ensure_directory_exists(optim_checkpoint_name) if not optimizer.is_stub_optimizer: optimizer.save_parameter_state(optim_checkpoint_name) + + # LayerWiseDistributedOptimizer save + if getattr(args, "optimizer", "adam").startswith("dist_") and args.ckpt_format == 'torch': + dp_rank = mpu.get_data_parallel_rank() + optim_checkpoint_name = os.path.join(os.path.dirname(checkpoint_name), f"layer_wise_optimizer_{dp_rank}.pt") + ensure_directory_exists(optim_checkpoint_name) + if not optimizer.is_stub_optimizer: + optimizer.save_state_dict_to_file(optim_checkpoint_name) async_save_request = None if args.async_save: @@ -1653,7 +1661,11 @@ def load_model_state_dict(module, state_dict, strict: bool): if not release and not args.finetune and not args.no_load_optim: try: # Load state dict. - if not skip_load_to_model_and_opt and optimizer is not None and not optimizer.is_stub_optimizer: + if getattr(args, "optimizer", "adam").startswith("dist_") and args.ckpt_format == 'torch': + dp_rank = mpu.get_data_parallel_rank() + optim_checkpoint_name = os.path.join(os.path.dirname(checkpoint_name), f"layer_wise_optimizer_{dp_rank}.pt") + optimizer.load_state_dict_from_file(optim_checkpoint_name) + elif not skip_load_to_model_and_opt and optimizer is not None and not optimizer.is_stub_optimizer: optimizer.load_state_dict(state_dict['optimizer']) # Load distributed optimizer's custom parameter state. From b39d10ab0186375141844e0cc7163b0226b75baf Mon Sep 17 00:00:00 2001 From: Boxiang Wang Date: Thu, 13 Nov 2025 16:34:43 -0800 Subject: [PATCH 3/4] Revert some changes Signed-off-by: Boxiang Wang --- .../core/extensions/transformer_engine.py | 23 ------------------- megatron/core/optimizer/muon.py | 2 -- megatron/training/arguments.py | 13 ----------- 3 files changed, 38 deletions(-) diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 33aa4b76b58..763bd6365b2 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -990,14 +990,6 @@ def __init__( self.kept_packed_seq_params.discard("cu_seqlens_q_padded") self.kept_packed_seq_params.discard("cu_seqlens_kv_padded") - if config.qk_clip or config.log_max_attention_logit: - # qk-clip is only supported in TE 2.9.0 and later - # assert is_te_min_version("2.9.0"), "qk-clip is only supported in TE 2.9.0 and later" - - # TE 2.9.0 introduces return_max_logit for qk-clip getting the max attention logits - extra_kwargs["return_max_logit"] = True - self.current_max_attn_logits = None - super().__init__( num_attention_heads=self.config.num_attention_heads, kv_channels=kv_channels, @@ -1068,21 +1060,6 @@ def forward( **packed_seq_kwargs, ) - if self.config.qk_clip or self.config.log_max_attention_logit: - # qk-clip is only supported in TE 2.9.0 and later - # assert is_te_min_version("2.9.0"), "qk-clip is only supported in TE 2.9.0 and later" - - # Update Q K outside of TE Attention API - core_attn_out, batch_max_attention_logits = core_attn_out - - # Update QK_Clip balancing eta - if self.current_max_attn_logits is None: - self.current_max_attn_logits = batch_max_attention_logits - else: - self.current_max_attn_logits = torch.max( - self.current_max_attn_logits, batch_max_attention_logits - ) - else: core_attn_out = super().forward( query, key, value, attention_mask, **attention_bias_kwargs, **packed_seq_kwargs diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py index 9fe132ea1d3..a31c84a6e8a 100644 --- a/megatron/core/optimizer/muon.py +++ b/megatron/core/optimizer/muon.py @@ -50,7 +50,6 @@ def __init__( use_nesterov: bool = True, weight_decay: float = 0.01, use_decoupled_weight_decay: bool = True, - use_independent_wd: bool = False, split_qkv: bool = False, is_qkv_fn: Callable[[torch.Tensor], bool] | None = None, qkv_split_shapes: tuple[int, int, int] | None = None, @@ -103,7 +102,6 @@ def scaled_orthogonalize_fn( use_nesterov, weight_decay, use_decoupled_weight_decay, - use_independent_wd, fp32_matmul_prec, scaled_orthogonalize_fn, ) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 4c7847f3927..3056f2007f2 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1003,19 +1003,6 @@ def validate_args(args, defaults={}): if args.add_bias_linear: args.add_qkv_bias = True - if args.qk_clip: - # assert is_te_min_version("2.9.0"), \ - # '--qk-clip is only supported with TE >= 2.9.0.' - assert 0.0 < args.qk_clip_alpha < 1.0, \ - '--qk-clip-alpha must be between 0.0 and 1.0 when using --qk-clip.' - assert args.qk_clip_threshold > 0, \ - '--qk-clip-threshold must be greater than 0 when using --qk-clip.' - - # decoupled log max attention logit check - # if args.log_max_attention_logit: - # assert is_te_min_version("2.9.0"), \ - # '--log-max-attention-logit is only supported with TE >= 2.9.0.' - # Retro checks. if args.retro_add_retriever: From 66030fd3efabd363c8b1b2123b0bd786c1a16d90 Mon Sep 17 00:00:00 2001 From: Boxiang Wang Date: Thu, 13 Nov 2025 17:23:31 -0800 Subject: [PATCH 4/4] Fix Signed-off-by: Boxiang Wang --- megatron/core/optimizer/optimizer.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index f3e256f6da7..51425de0156 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -1224,10 +1224,9 @@ def load_state_dict(self, state_dict): for optimizer, state in zip(self.chained_optimizers, state_dict): # Mostly for Layerwise Dist Opt situation but it can be a general check # if the state dict optimizer is in this rank, load the state dict - maybe_optimizer_state = state['optimizer'] if 'optimizer' in state else state['optimizer_state_dict'] - if 'state' in maybe_optimizer_state: - pass - # optimizer.load_state_dict(state) + if 'optimizer' in state: + if 'state' in state['optimizer']: + optimizer.load_state_dict(state) # else: # print(f"state_dict: {torch.distributed.get_rank()} {state['optimizer']} {type(optimizer)} \n") # assert False