diff --git a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py b/megatron/core/tensor_parallel/generalized_tensor_parallelism.py index 5eac911d1b3..5d95347ef87 100644 --- a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py +++ b/megatron/core/tensor_parallel/generalized_tensor_parallelism.py @@ -217,6 +217,10 @@ def _classify_param_chain(param_name: str) -> str: if not scope: # CG disabled return U + # MoE latent projections. + if ".mlp.fc1_latent_proj." in n or ".mlp.fc2_latent_proj." in n: + return G if "moe_router" in scope else U + if ".mlp.shared_experts." in n: if _MOE_SHARED_EXPERT_OVERLAP: return U diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py index 297c89578f1..6c594d0be62 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py @@ -502,6 +502,42 @@ def test_graphness_helpers(self): assert gtp_module._chain_is_graphed("GTP_graphed") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="GTPShardedParam requires CUDA") +class TestLatentProjectionChainClassification: + """Classify opted-in latent projections through the public GTP chain API.""" + + FC1 = "decoder.layers.3.mlp.fc1_latent_proj.weight" + FC2 = "decoder.layers.3.mlp.fc2_latent_proj.weight" + + def teardown_method(self, method): + gtp_module.set_cuda_graph_modules(None, cuda_graph_impl="none") + gtp_module.reset_gtp_state() + + def _chains(self, *, cuda_graph_modules=None, cuda_graph_impl="none"): + params = tuple(GTPShardedParam(torch.zeros(1, device="cuda")) for _ in range(2)) + assert all(isinstance(param, GTPShardedParam) for param in params) + + class _Model: + def named_parameters(_self): + return iter(zip((self.FC1, self.FC2), params)) + + gtp_module.classify_gtp_remat_chains( + _Model(), cuda_graph_modules=cuda_graph_modules, cuda_graph_impl=cuda_graph_impl + ) + return tuple(param.chain_id for param in params) + + def test_eager_latent_projections_are_ungraphed(self): + assert self._chains() == (GTPChain.UNGRAPHED.value, GTPChain.UNGRAPHED.value) + + def test_local_router_captures_both_latent_projections(self): + chains = self._chains(cuda_graph_modules={"moe_router"}, cuda_graph_impl="local") + assert chains == (GTPChain.GRAPHED.value, GTPChain.GRAPHED.value) + + def test_unrelated_local_scope_leaves_latent_projections_ungraphed(self): + chains = self._chains(cuda_graph_modules={"mamba", "attn"}, cuda_graph_impl="local") + assert chains == (GTPChain.UNGRAPHED.value, GTPChain.UNGRAPHED.value) + + class TestGroupedDoubleBuffer: """One-block-ahead grouped chains must double-buffer: consecutive MoE layers get distinct gather buffers (else prefetching layer N+1 clobbers layer N's in-use weight). Pure cache-key diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_partial_cg.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_partial_cg.py index addcb9cb657..14641c3a3b4 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_gtp_partial_cg.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_partial_cg.py @@ -36,9 +36,10 @@ ) -def _worker_gtp_partial_cg_correctness(rank, world_size, port): - """Compare eager and local attention CUDA graphs with GTP2 x DP2.""" +def _worker_gtp_partial_cg_correctness(rank, world_size, port, partial_cg_modules, opt_in_modules): + """Compare eager and local CUDA graphs with GTP2 x DP2.""" del port + gtp_module._GTP_PARAMS.clear() from megatron.core import parallel_state as ps from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec @@ -56,14 +57,15 @@ def _worker_gtp_partial_cg_correctness(rank, world_size, port): ) from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp from megatron.core.transformer.transformer_config import TransformerConfig - - hidden = 4096 - num_heads = 32 - ffn_hidden = 16384 - # Four layers force parameters with matching scheduling domains/shapes to reuse the two-slot - # wgrad ring across independently replayed graphs. - num_layers = 4 - sequence_length = 32 + from megatron.core.transformer.transformer_layer import MoETransformerLayer + + latent_projection_case = "moe_latent_proj" in opt_in_modules + hidden = 256 if latent_projection_case else 4096 + num_heads = 8 if latent_projection_case else 32 + ffn_hidden = 512 if latent_projection_case else 16384 + # Use multiple layers to exercise repeated local CUDA-graph execution with GTP parameters. + num_layers = 1 if latent_projection_case else 4 + sequence_length = 16 if latent_projection_case else 32 batch_size = 1 learning_rate = 0.01 steps = 10 @@ -73,6 +75,18 @@ def _worker_gtp_partial_cg_correctness(rank, world_size, port): assert world_size == gtp_degree * dp_degree def make_config(*, partial_cg=False): + moe_options = {} + if latent_projection_case: + moe_options = { + "num_moe_experts": 2, + "moe_router_topk": 1, + "moe_router_pre_softmax": True, + "moe_ffn_hidden_size": ffn_hidden, + "moe_grouped_gemm": True, + "moe_token_dispatcher_type": "allgather", + "moe_aux_loss_coeff": 0.0, + "moe_latent_size": 128, + } return TransformerConfig( num_attention_heads=num_heads, num_layers=num_layers, @@ -83,16 +97,40 @@ def make_config(*, partial_cg=False): hidden_dropout=0.0, attention_dropout=0.0, bias_dropout_fusion=False, + gradient_accumulation_fusion=latent_projection_case, tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_weight_remat_size=gtp_degree, + gtp_remat_opt_in_modules=opt_in_modules, cuda_graph_impl="local" if partial_cg else "none", - cuda_graph_modules=["attn"] if partial_cg else [], + cuda_graph_modules=partial_cg_modules if partial_cg else [], cuda_graph_warmup_steps=2, + **moe_options, ) - def make_attention_stack(config, pg_collection): - spec = copy.deepcopy(get_gpt_layer_with_transformer_engine_spec()) + def make_layer_stack(config, pg_collection): + spec = copy.deepcopy( + get_gpt_layer_with_transformer_engine_spec( + num_experts=2 if latent_projection_case else None, + moe_grouped_gemm=latent_projection_case, + ) + ) + if latent_projection_case: + spec.submodules.input_layernorm = IdentityOp + spec.submodules.self_attention = IdentityOp + spec.submodules.self_attn_bda = IdentityFuncOp + return torch.nn.ModuleList( + [ + MoETransformerLayer( + config, + spec.submodules, + layer_number=1, + pg_collection=pg_collection, + name="decoder.layers.0", + ) + ] + ) + spec.submodules.pre_mlp_layernorm = IdentityOp spec.submodules.mlp = IdentityOp spec.submodules.mlp_bda = IdentityFuncOp @@ -105,6 +143,27 @@ def make_attention_stack(config, pg_collection): ] ) + def get_cudagraph_managers(layers): + if latent_projection_case: + return [ + manager + for layer in layers + for manager in (layer.cudagraph_manager_router, layer.cudagraph_manager_postprocess) + ] + return [layer.cudagraph_manager for layer in layers] + + def get_latent_params(layers): + return [ + param + for name, param in layers.named_parameters() + if "fc1_latent_proj.weight" in name or "fc2_latent_proj.weight" in name + ] + + def make_pg_collection(): + if latent_projection_case: + return ProcessGroupCollection.use_mpu_process_groups() + return ProcessGroupCollection.use_mpu_process_groups(required_pgs=["tp", "cp", "gtp_remat"]) + def run_step(layers, x): with fp8_autocast(enabled=False): for layer in layers: @@ -162,17 +221,19 @@ def apply_sgd_step(layers, gtp_size): tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=gtp_degree ) model_parallel_cuda_manual_seed(42) - pg_collection = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=["tp", "cp", "gtp_remat"] - ) + pg_collection = make_pg_collection() eager_config = make_config() - eager = make_attention_stack(eager_config, pg_collection).cuda() + eager = make_layer_stack(eager_config, pg_collection).cuda() eager_gtp_group = ps.get_gtp_weight_remat_group() eager_dp_group = ps.get_data_parallel_group(with_gtp_remat=False) eager_dp_rank = eager_dp_group.rank() assert eager_gtp_group.size() == gtp_degree assert eager_dp_group.size() == dp_degree assert any(isinstance(param, GTPShardedParam) for param in eager.parameters()) + if latent_projection_case: + eager_latent_params = get_latent_params(eager) + assert len(eager_latent_params) == 2 + assert all(isinstance(param, GTPShardedParam) for param in eager_latent_params) initialize_main_grads(eager) saved_local_weights = {name: param.data.clone() for name, param in eager.named_parameters()} @@ -190,20 +251,20 @@ def apply_sgd_step(layers, gtp_size): apply_sgd_step(eager, eager_gtp_group.size()) del eager, loss, x + torch.cuda.synchronize() ps.destroy_model_parallel() gtp_module.reset_gtp_state() + gtp_module._GTP_PARAMS.clear() - # Optimized path: the same GTP2 x DP2 topology with attention-only local CUDA graphs. + # Optimized path: the same GTP2 x DP2 topology with local CUDA graphs. ps.initialize_model_parallel( tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=gtp_degree ) initialize_rng_tracker(use_te_rng_tracker=True, force_reset=True) model_parallel_cuda_manual_seed(42) - pg_collection = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=["tp", "cp", "gtp_remat"] - ) + pg_collection = make_pg_collection() partial_cg_config = make_config(partial_cg=True) - partial_cg = make_attention_stack(partial_cg_config, pg_collection).cuda() + partial_cg = make_layer_stack(partial_cg_config, pg_collection).cuda() classify_gtp_remat_chains( partial_cg, cuda_graph_modules=partial_cg_config.cuda_graph_modules, @@ -220,7 +281,8 @@ def apply_sgd_step(layers, gtp_size): assert dp_rank == eager_dp_rank gtp_params = [param for param in partial_cg.parameters() if isinstance(param, GTPShardedParam)] assert gtp_params, "GTP not active: no GTPShardedParam found" - assert all(param.chain_id == GTPChain.GRAPHED.value for param in gtp_params) + params_in_scope = get_latent_params(partial_cg) if latent_projection_case else gtp_params + assert all(param.chain_id == GTPChain.GRAPHED.value for param in params_in_scope) for name, param in partial_cg.named_parameters(): param.data.copy_(saved_local_weights[name]) # Production captures after DDP maps every parameter into a main-grad buffer and initializes @@ -241,18 +303,21 @@ def apply_sgd_step(layers, gtp_size): wait_for_gtp_grad_reduction_on_current_stream() eager_grad_norm = global_grad_norm(partial_cg, gtp_group) eager_probe_loss_value = eager_probe_loss.item() + del eager_probe_loss, eager_probe_x + reset_grad_state(partial_cg) create_cudagraphs() assert _CudagraphGlobalRecord.cudagraph_created - runners = [layer.cudagraph_manager.cudagraph_runners[0] for layer in partial_cg] - assert all(runner.gtp_remat for runner in runners) - assert any(runner._gtp_wgrad_ring_slots for runner in runners) + managers = get_cudagraph_managers(partial_cg) + assert all(len(manager.cudagraph_runners) == 1 for manager in managers) + runners = [manager.cudagraph_runners[0] for manager in managers] + assert any(runner.gtp_remat for runner in runners) replay_grad_norms = [] replay_losses = [] for _ in range(3): reset_grad_state(partial_cg) - replay_x = eager_probe_x.detach().clone().requires_grad_() + replay_x = make_replica_input(1234, dp_rank).requires_grad_() replay_loss = run_step(partial_cg, replay_x) replay_loss.backward() wait_for_gtp_grad_reduction_on_current_stream() @@ -280,7 +345,7 @@ def apply_sgd_step(layers, gtp_size): flush=True, ) - del eager_probe_loss, eager_probe_x, replay_loss, replay_x + del replay_loss, replay_x for step in range(steps): reset_grad_state(partial_cg) @@ -295,19 +360,21 @@ def apply_sgd_step(layers, gtp_size): del loss, x finally: torch.cuda.synchronize() - for layer in partial_cg: - for runner in layer.cudagraph_manager.cudagraph_runners: + managers = get_cudagraph_managers(partial_cg) + for manager in managers: + for runner in manager.cudagraph_runners: if runner.fwd_graph is not None: runner.fwd_graph.reset() if runner.bwd_graph is not None: runner.bwd_graph.reset() delete_cuda_graphs() - for layer in partial_cg: - layer.cudagraph_manager.cudagraph_runners.clear() + for manager in managers: + manager.cudagraph_runners.clear() gc.collect() ps.destroy_model_parallel() ps.initialize_model_parallel() gtp_module.reset_gtp_state() + gtp_module._GTP_PARAMS.clear() if rank == 0: for step, (eager_loss, partial_cg_loss) in enumerate(zip(eager_losses, partial_cg_losses)): @@ -324,8 +391,17 @@ def apply_sgd_step(layers, gtp_size): class TestGTPPartialCGCorrectness: - def test_gtp_partial_cg_loss_and_grad_norm_match_eager(self): + @pytest.mark.parametrize( + "partial_cg_modules,opt_in_modules", + [ + pytest.param(["attn"], [], id="attention"), + pytest.param(["moe_router"], ["moe_latent_proj"], id="moe-router-latent-projections"), + ], + ) + def test_gtp_partial_cg_loss_and_grad_norm_match_eager( + self, partial_cg_modules, opt_in_modules + ): """Local-CG loss trajectory and global grad norm must match eager execution.""" if torch.cuda.device_count() < 4: pytest.skip("Requires at least 4 CUDA devices") - _run_distributed(_worker_gtp_partial_cg_correctness, 4) + _run_distributed(_worker_gtp_partial_cg_correctness, 4, partial_cg_modules, opt_in_modules)