diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py index 53522dd8b2b..10b3db21ca1 100644 --- a/megatron/core/models/common/language_module/language_module.py +++ b/megatron/core/models/common/language_module/language_module.py @@ -6,7 +6,7 @@ import torch from torch import Tensor -from megatron.core import parallel_state, tensor_parallel +from megatron.core import tensor_parallel from megatron.core.dist_checkpointing.mapping import ShardedStateDict from megatron.core.transformer.cuda_graphs import CudaGraphManager @@ -28,6 +28,7 @@ from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.utils import ensure_metadata_has_dp_cp_group from megatron.core.utils import ( + get_pg_rank, get_tensor_model_parallel_group_if_none, is_te_min_version, make_tp_sharded_tensor_for_checkpoint, @@ -509,7 +510,7 @@ def tie_embeddings_and_output_weights_state_dict( last_stage_word_emb_replica_id = ( 1, # copy of first stage embedding 0, - parallel_state.get_data_parallel_rank(with_context_parallel=True), + get_pg_rank(metadata['dp_cp_group']), ) sharded_state_dict[output_layer_weight_key] = make_tp_sharded_tensor_for_checkpoint( diff --git a/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py b/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py index 10d9d67c151..24232ebe8df 100644 --- a/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py +++ b/tests/unit_tests/models/mimo/test_mimo_1f1b_schedule.py @@ -213,6 +213,7 @@ def get_language_model_spec( bias=True, dropout=True, per_token_loss=False, + share_embeddings_and_output_weights=False, ): """Get the language model spec. @@ -264,6 +265,7 @@ def get_language_model_spec( "max_sequence_length": seq_len, "pre_process": (pp_rank == 0), "post_process": (pp_rank == pp_size - 1), + "share_embeddings_and_output_weights": share_embeddings_and_output_weights, "pg_collection": pg_collection, }, ) @@ -382,6 +384,7 @@ def get_mimo_model( dropout=True, per_token_loss=False, use_layer_wise_distributed_optimizer=False, + share_embeddings_and_output_weights=False, ): """Create MIMO model with TransformerBlock encoder and GPTModel LLM. @@ -403,6 +406,8 @@ def get_mimo_model( and LLM without relying on the per-DDP built-in scaling. use_layer_wise_distributed_optimizer: Whether to wrap active modules through the production MIMO LayerWise parameter-layout path. + share_embeddings_and_output_weights: If True, tie the LLM word embedding and + output-layer weights (GPTModel kwarg of the same name). """ language_pg = get_pg_collection_with_embedding_groups(llm_grid, is_language_model=True) vision_pg = get_pg_collection_with_embedding_groups(encoder_grid, is_language_model=False) @@ -418,6 +423,7 @@ def get_mimo_model( bias=bias, dropout=dropout, per_token_loss=per_token_loss, + share_embeddings_and_output_weights=share_embeddings_and_output_weights, ) vision_submodule_spec = get_vision_submodules_spec( num_layers=num_layers, diff --git a/tests/unit_tests/models/mimo/test_mimo_checkpoint.py b/tests/unit_tests/models/mimo/test_mimo_checkpoint.py index dc12e53571a..5779d3b695a 100644 --- a/tests/unit_tests/models/mimo/test_mimo_checkpoint.py +++ b/tests/unit_tests/models/mimo/test_mimo_checkpoint.py @@ -56,7 +56,49 @@ def _randomize_params(model, seed): p.random_() -def _create_model_and_optimizer(encoder_grid, llm_grid, hidden_size, num_layers, vocab_size, seed): +def _sync_tied_pair(model, first_rank, last_rank): + """Copy the first-stage word embedding onto the last-stage tied output layer. + + Checkpointing tied embeddings assumes both copies hold identical values, as training + keeps them via the embedding-group grad all-reduce. _randomize_params breaks that + invariant, so the tied test restores it before the fp32 main params are cloned at + optimizer construction. All ranks must call this (collective broadcast). The fake + per-param grads do not need syncing: the optimizer reads the zero-initialized + DDP main_grad buffers, not param.grad, so the step is grad-free either way. + """ + rank = dist.get_rank() + meta = [None] + src = None + if rank == first_rank: + emb = next( + p + for name, p in model.named_parameters() + if name.endswith('embedding.word_embeddings.weight') + ) + src = emb.data + meta = [(tuple(src.shape), src.dtype)] + dist.broadcast_object_list(meta, src=first_rank) + shape, dtype = meta[0] + buf = src.contiguous() if rank == first_rank else torch.empty(shape, dtype=dtype, device='cuda') + dist.broadcast(buf, src=first_rank) + if rank == last_rank: + out = next( + p for name, p in model.named_parameters() if name.endswith('output_layer.weight') + ) + with torch.no_grad(): + out.data.copy_(buf) + + +def _create_model_and_optimizer( + encoder_grid, + llm_grid, + hidden_size, + num_layers, + vocab_size, + seed, + tie_embeddings=False, + tied_sync_ranks=None, +): """Create MIMO model with DDP + optimizer, do a fake step to populate optimizer state. Caller must call create_all_embedding_groups() before this function. @@ -71,8 +113,11 @@ def _create_model_and_optimizer(encoder_grid, llm_grid, hidden_size, num_layers, num_layers=num_layers, vocab_size=vocab_size, seq_len=64, + share_embeddings_and_output_weights=tie_embeddings, ) _randomize_params(mimo_model, seed) + if tied_sync_ranks is not None: + _sync_tied_pair(mimo_model, *tied_sync_ranks) # Use Float16Optimizer (not DistributedOptimizer) to exercise the MIMO-specific # param_groups/grad_scaler extraction in sharded_state_dict. DistributedOptimizer @@ -107,6 +152,7 @@ def run_checkpoint_test( hidden_size=256, num_layers=2, vocab_size=1000, + tie_embeddings=False, ): """Save model + optimizer checkpoint, load into fresh instances, verify match.""" # Clear NVTE env vars that the conftest set_env fixture sets to '0'. @@ -121,9 +167,23 @@ def run_checkpoint_test( llm_grid = create_hypercomm_grid(offset=llm_offset, tp=llm_tp, cp=1, pp=llm_pp, dp=llm_dp) create_all_embedding_groups([encoder_grid, llm_grid]) + tied_sync_ranks = None + if tie_embeddings: + # The tied pair lives on the first and last LLM PP stage; the sync helper + # assumes each stage is a single rank. + assert llm_tp == 1 and llm_dp == 1, "tie_embeddings test path assumes llm tp=1, dp=1" + tied_sync_ranks = (llm_offset, llm_offset + llm_pp - 1) + # --- Create model A + optimizer, snapshot state --- model_a, optimizer_a = _create_model_and_optimizer( - encoder_grid, llm_grid, hidden_size, num_layers, vocab_size, seed=1 + encoder_grid, + llm_grid, + hidden_size, + num_layers, + vocab_size, + seed=1, + tie_embeddings=tie_embeddings, + tied_sync_ranks=tied_sync_ranks, ) params_a = {name: p.clone() for name, p in model_a.named_parameters()} @@ -150,7 +210,14 @@ def run_checkpoint_test( # --- Create model B + optimizer with different weights (reuse same grids) --- model_b, optimizer_b = _create_model_and_optimizer( - encoder_grid, llm_grid, hidden_size, num_layers, vocab_size, seed=2 + encoder_grid, + llm_grid, + hidden_size, + num_layers, + vocab_size, + seed=2, + tie_embeddings=tie_embeddings, + tied_sync_ranks=tied_sync_ranks, ) # Load model @@ -259,6 +326,26 @@ def test_encoder_tp1_llm_pp7(self): num_layers=7, ) + def test_encoder_tp1_llm_pp7_tied_embeddings(self): + """Tied word embeddings with PP >= 2: saving reaches the tied output-layer + replica_id in tie_embeddings_and_output_weights_state_dict, which must come + from metadata['dp_cp_group'] — the global MPU is not initialized here.""" + if self.world_size != 8: + pytest.skip(f"Requires 8 GPUs, got {self.world_size}") + run_checkpoint_test( + encoder_tp=1, + encoder_pp=1, + encoder_dp=1, + encoder_offset=0, + llm_tp=1, + llm_pp=7, + llm_dp=1, + llm_offset=1, + hidden_size=256, + num_layers=7, + tie_embeddings=True, + ) + def test_encoder_tp2_pp2_llm_tp2_pp2(self): if self.world_size != 8: pytest.skip(f"Requires 8 GPUs, got {self.world_size}")