diff --git a/docs/api-guide/core/generalized_tensor_parallel.md b/docs/api-guide/core/generalized_tensor_parallel.md index a14475cd869..9c4dcda51df 100644 --- a/docs/api-guide/core/generalized_tensor_parallel.md +++ b/docs/api-guide/core/generalized_tensor_parallel.md @@ -557,5 +557,6 @@ torchrun --nproc-per-node 4 -m pytest tests/unit_tests/generalized_tensor_parall | `test_gtp_dcp.py` | DCP sharding metadata (§3.3): TP×GTP_remat offsets, pad reshard, `replica_id`, native-FP8 save/load. | | `test_gtp_muon_dcp.py` | Muon optimizer-state DCP roundtrip (§1.6): `replica_id` fold + native-FP8 backfill matching. | | `test_gtp_fp8_param_gather.py` | Native-FP8 GTP_remat (§1.3): fp8-vs-BF16 loss parity (TP1/TP2, MoE), post-save-spike guard. | +| `test_gtp_custom_pgs.py` | `pg_collection` plumbing: a custom `gtp_remat` group (permuted ranks, same size) must give the same fwd/bwd results as the MPU groups — catches modules reading `parallel_state` instead of the collection passed to them. | All tests require ≥ 4 GPUs and TransformerEngine >= 2.19; they self-skip when those are unavailable. A green run (skips for unmet hardware/config are acceptable) is the minimum bar for any GTP_remat change. diff --git a/examples/mimo/model_providers/nemotron_moe_vlm.py b/examples/mimo/model_providers/nemotron_moe_vlm.py index 133bf9bc1d2..03013e73e85 100644 --- a/examples/mimo/model_providers/nemotron_moe_vlm.py +++ b/examples/mimo/model_providers/nemotron_moe_vlm.py @@ -12,6 +12,7 @@ from examples.mimo.model_providers.radio_encoder import ( RADIO_ENCODER_MODULE_NAME, _base_config, + _disable_gtp, _make_dense_non_hybrid, add_radio_encoder_args, radio_vision_config, @@ -21,6 +22,7 @@ from megatron.core.activations import squared_relu from megatron.core.hyper_comm_grid import HyperCommGrid from megatron.core.hyper_comm_grid import _is_process_group_member as is_process_group_member +from megatron.core.model_parallel_config import resolve_tensor_parallel_weight_shards from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec from megatron.core.models.mamba.mamba_model import MambaModel from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY @@ -105,6 +107,23 @@ def nemotron_language_config( config.expert_tensor_parallel_size = expt_tp_size config.tensor_model_parallel_size = tp_size config.pipeline_model_parallel_size = pp_size + config.tensor_parallel_num_weight_shards, config.gtp_weight_remat_size = ( + resolve_tensor_parallel_weight_shards( + tp_size, + getattr(args, "tensor_parallel_num_weight_shards", None), + getattr(args, "gtp_weight_remat_size", 1), + ) + ) + ( + config.expert_tensor_parallel_num_weight_shards, + config.expert_gtp_weight_remat_size, + ) = resolve_tensor_parallel_weight_shards( + expt_tp_size, + getattr(args, "expert_tensor_parallel_num_weight_shards", None), + getattr(args, "expert_gtp_weight_remat_size", 1), + shards_field="expert_tensor_parallel_num_weight_shards", + tp_field="expert_tensor_parallel_size", + ) config.sequence_parallel = tp_size > 1 config.position_embedding_type = "none" return config @@ -142,6 +161,7 @@ def nemotron_projection_config( config.normalization = "RMSNorm" _make_dense_non_hybrid(config) # Projection inherits no MoE/Mamba/hybrid settings. config.tensor_model_parallel_size = tp_size + _disable_gtp(config) config.sequence_parallel = False return config diff --git a/examples/mimo/model_providers/radio_encoder.py b/examples/mimo/model_providers/radio_encoder.py index 9e0591cc7e7..c301c0c29ca 100644 --- a/examples/mimo/model_providers/radio_encoder.py +++ b/examples/mimo/model_providers/radio_encoder.py @@ -83,6 +83,15 @@ def _make_dense_non_hybrid(config: TransformerConfig) -> None: config.use_fused_weighted_squared_relu = False +def _disable_gtp(config: TransformerConfig) -> None: + """Keep encoder-side weights replicated across the language module's GTP axes.""" + config.tensor_parallel_num_weight_shards = config.tensor_model_parallel_size + config.gtp_weight_remat_size = 1 + expert_tp = config.expert_tensor_parallel_size or config.tensor_model_parallel_size + config.expert_tensor_parallel_num_weight_shards = expert_tp + config.expert_gtp_weight_remat_size = 1 + + def radio_vision_config(args: argparse.Namespace, tp_size: int, pp_size: int) -> TransformerConfig: """RADIO vision config: stock from-args base + RADIO-specific overrides.""" config = deepcopy(_base_config(args)) @@ -114,6 +123,7 @@ def radio_vision_config(args: argparse.Namespace, tp_size: int, pp_size: int) -> config.bf16 = bf16 config.tensor_model_parallel_size = tp_size config.pipeline_model_parallel_size = pp_size + _disable_gtp(config) config.sequence_parallel = False return config diff --git a/examples/mimo/pretrain_mimo.py b/examples/mimo/pretrain_mimo.py index ee521b188a5..06219a286fb 100644 --- a/examples/mimo/pretrain_mimo.py +++ b/examples/mimo/pretrain_mimo.py @@ -48,14 +48,19 @@ def _parse_and_validate() -> argparse.Namespace: args = parse_args(extra_args_provider) validate_hetero_grid_args(args, args.world_size) physical_world_size = args.world_size - # Stock validate_args sets data_parallel_size = world_size // (tp*pp*cp); feed the - # language module's world (llm_dp; stock tp/pp/cp stay 1, MIMO parallelism is in --llm-*) - # so it yields llm_dp. The physical world incl. encoder ranks is restored below. + # Stock validation owns the derived training arguments. Give it the language module's + # parallel degrees and logical world so its DP/GTP accounting matches the explicit MIMO grid. + args.tensor_model_parallel_size = args.llm_tp + args.pipeline_model_parallel_size = args.llm_pp + args.context_parallel_size = args.llm_cp + args.expert_model_parallel_size = args.llm_ep + args.expert_tensor_parallel_size = args.llm_expt_tp or 1 args.world_size = ( args.llm_dp - * args.tensor_model_parallel_size - * args.pipeline_model_parallel_size - * args.context_parallel_size + * args.llm_tp + * args.gtp_weight_remat_size + * args.llm_pp + * args.llm_cp ) try: validate_args(args, {"dataloader_type": "external"}) diff --git a/examples/mimo/scripts/run_hetero_nemotron_20l_mock_train.sh b/examples/mimo/scripts/run_hetero_nemotron_20l_mock_train.sh index 787df50fef8..92bdce51a25 100755 --- a/examples/mimo/scripts/run_hetero_nemotron_20l_mock_train.sh +++ b/examples/mimo/scripts/run_hetero_nemotron_20l_mock_train.sh @@ -1,26 +1,78 @@ #!/bin/bash -# Run an eight-rank heterogeneous mock training loop with Nemotron6-MoE VLM 20L. +# Run heterogeneous mock training with the Nemotron6-MoE VLM 20L recipe. set -euo pipefail +export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1} + TRAIN_ITERS=${TRAIN_ITERS:-20} NUM_MICROBATCHES=${NUM_MICROBATCHES:-4} EVAL_INTERVAL=${EVAL_INTERVAL:-1} EVAL_ITERS=${EVAL_ITERS:-0} -MICRO_BATCH_SIZE=1 -LLM_DP=2 -GLOBAL_BATCH_SIZE=$((MICRO_BATCH_SIZE * NUM_MICROBATCHES * LLM_DP)) +MICRO_BATCH_SIZE=${MICRO_BATCH_SIZE:-1} +NNODES=${NNODES:-1} +NPROC_PER_NODE=${NPROC_PER_NODE:-8} +ENCODER_TP=${ENCODER_TP:-2} +ENCODER_DP=${ENCODER_DP:-2} +LLM_OFFSET=${LLM_OFFSET:-$((ENCODER_TP * ENCODER_DP))} +LLM_TP=${LLM_TP:-2} +LLM_CP=${LLM_CP:-1} +LLM_PP=${LLM_PP:-1} +LLM_DP=${LLM_DP:-2} +LLM_EP=${LLM_EP:-4} +LLM_EXPT_TP=${LLM_EXPT_TP:-1} +TENSOR_PARALLEL_NUM_WEIGHT_SHARDS=${TENSOR_PARALLEL_NUM_WEIGHT_SHARDS:-${LLM_TP}} +EXPERT_TENSOR_PARALLEL_NUM_WEIGHT_SHARDS=${EXPERT_TENSOR_PARALLEL_NUM_WEIGHT_SHARDS:-${LLM_EXPT_TP}} + +if ((TENSOR_PARALLEL_NUM_WEIGHT_SHARDS % LLM_TP != 0)); then + echo "TENSOR_PARALLEL_NUM_WEIGHT_SHARDS must be divisible by LLM_TP" >&2 + exit 2 +fi +if ((EXPERT_TENSOR_PARALLEL_NUM_WEIGHT_SHARDS % LLM_EXPT_TP != 0)); then + echo "EXPERT_TENSOR_PARALLEL_NUM_WEIGHT_SHARDS must be divisible by LLM_EXPT_TP" >&2 + exit 2 +fi + +GTP=$((TENSOR_PARALLEL_NUM_WEIGHT_SHARDS / LLM_TP)) +LLM_SIZE=$((LLM_TP * GTP * LLM_CP * LLM_PP * LLM_DP)) +EXPECTED_WORLD_SIZE=$((ENCODER_TP * ENCODER_DP + LLM_SIZE)) +WORLD_SIZE=$((NNODES * NPROC_PER_NODE)) +if ((WORLD_SIZE != EXPECTED_WORLD_SIZE)); then + echo "NNODES*NPROC_PER_NODE=${WORLD_SIZE}, but encoder+LLM grids require ${EXPECTED_WORLD_SIZE}" >&2 + exit 2 +fi + +GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE:-$((MICRO_BATCH_SIZE * NUM_MICROBATCHES * LLM_DP * GTP))} TORCHRUN_LOG_DIR=${TORCHRUN_LOG_DIR:-"${PWD}/logs/torchrun-$(date +%Y%m%d_%H%M%S)-$$"} mkdir -p "${TORCHRUN_LOG_DIR}" TORCHRUN_ARGS=( - --standalone - --nproc-per-node 8 + --nproc-per-node "${NPROC_PER_NODE}" --log-dir "${TORCHRUN_LOG_DIR}" --redirects 3 - --tee 3 + --tee 0:3 ) +if ((NNODES == 1)); then + TORCHRUN_ARGS+=(--standalone) +else + NODE_RANK=${NODE_RANK:-${SLURM_NODEID:-0}} + if [[ -z "${MASTER_ADDR:-}" ]]; then + if [[ -z "${SLURM_JOB_NODELIST:-}" ]]; then + echo "MASTER_ADDR or SLURM_JOB_NODELIST is required for a multi-node run" >&2 + exit 2 + fi + mapfile -t slurm_nodes < <(scontrol show hostnames "${SLURM_JOB_NODELIST}") + MASTER_ADDR=${slurm_nodes[0]} + fi + MASTER_PORT=${MASTER_PORT:-$((10000 + ${SLURM_JOB_ID:-0} % 50000))} + TORCHRUN_ARGS+=( + --nnodes "${NNODES}" + --node-rank "${NODE_RANK}" + --master-addr "${MASTER_ADDR}" + --master-port "${MASTER_PORT}" + ) +fi uv run --extra ssm python -m torch.distributed.run \ "${TORCHRUN_ARGS[@]}" \ @@ -71,15 +123,17 @@ uv run --extra ssm python -m torch.distributed.run \ --seq-length 8192 \ --max-position-embeddings 8192 \ --bf16 \ - --encoder-tp 2 \ - --encoder-dp 2 \ - --llm-offset 4 \ - --llm-tp 2 \ - --llm-cp 1 \ - --llm-pp 1 \ + --encoder-tp "${ENCODER_TP}" \ + --encoder-dp "${ENCODER_DP}" \ + --llm-offset "${LLM_OFFSET}" \ + --llm-tp "${LLM_TP}" \ + --llm-cp "${LLM_CP}" \ + --llm-pp "${LLM_PP}" \ --llm-dp "${LLM_DP}" \ - --llm-ep 4 \ - --llm-expt-tp 1 \ + --llm-ep "${LLM_EP}" \ + --llm-expt-tp "${LLM_EXPT_TP}" \ + --tensor-parallel-num-weight-shards "${TENSOR_PARALLEL_NUM_WEIGHT_SHARDS}" \ + --expert-tensor-parallel-num-weight-shards "${EXPERT_TENSOR_PARALLEL_NUM_WEIGHT_SHARDS}" \ --vocab-size 131072 \ --micro-batch-size "${MICRO_BATCH_SIZE}" \ --global-batch-size "${GLOBAL_BATCH_SIZE}" \ diff --git a/examples/mimo/training/args.py b/examples/mimo/training/args.py index e62bc31967b..ad666a7aef3 100644 --- a/examples/mimo/training/args.py +++ b/examples/mimo/training/args.py @@ -8,6 +8,7 @@ from typing import List from examples.mimo.training.topology import ModuleGridSpec +from megatron.core.model_parallel_config import resolve_tensor_parallel_weight_shards from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY @@ -61,6 +62,7 @@ def add_hetero_grid_args(parser: argparse.ArgumentParser) -> argparse.ArgumentPa def validate_hetero_grid_args(args: argparse.Namespace, world_size: int) -> tuple[int, int]: """Validate the disjoint hetero grid layout; returns ``(encoder_size, llm_size)``.""" + resolve_hetero_gtp_args(args) if args.llm_cp != 1: raise ValueError("hetero MIMO training currently supports CP=1 only") @@ -76,7 +78,7 @@ def validate_hetero_grid_args(args: argparse.Namespace, world_size: int) -> tupl f"--num-experts ({num_experts}) must be divisible by --llm-ep ({args.llm_ep})" ) - llm_size = args.llm_tp * args.llm_cp * args.llm_pp * args.llm_dp + llm_size = args.llm_tp * args.gtp_weight_remat_size * args.llm_cp * args.llm_pp * args.llm_dp if args.llm_only: if getattr(args, "encoder_ddp_overlap", False): @@ -94,12 +96,14 @@ def validate_hetero_grid_args(args: argparse.Namespace, world_size: int) -> tupl ) return 0, llm_size - # Fan-out divisibility: the bridge splits (mbs * llm_dp) LLM lanes across + # Fan-out divisibility: the bridge splits (mbs * llm_dp * gtp) LLM lanes across # encoder_dp encoder lanes; the split must be exact. - if (args.micro_batch_size * args.llm_dp) % args.encoder_dp != 0: + llm_data_parallel_size = args.llm_dp * args.gtp_weight_remat_size + if (args.micro_batch_size * llm_data_parallel_size) % args.encoder_dp != 0: raise ValueError( - "--micro-batch-size * --llm-dp must be divisible by --encoder-dp " - f"(got {args.micro_batch_size} * {args.llm_dp} % {args.encoder_dp} != 0)" + "--micro-batch-size * --llm-dp * GTP must be divisible by --encoder-dp " + f"(got {args.micro_batch_size} * {args.llm_dp} * " + f"{args.gtp_weight_remat_size} % {args.encoder_dp} != 0)" ) encoder_size = args.encoder_tp * args.encoder_dp @@ -134,8 +138,10 @@ def build_module_grid_specs( cp=args.llm_cp, pp=args.llm_pp, ep=args.llm_ep, + gtp_remat=args.gtp_weight_remat_size, rank_offset=args.llm_offset, expt_tp=args.llm_expt_tp or 1, + expt_gtp_remat=args.expert_gtp_weight_remat_size, ) if args.llm_only: @@ -154,6 +160,27 @@ def build_module_grid_specs( return [encoder_grid_spec, language_grid_spec] +def resolve_hetero_gtp_args(args: argparse.Namespace) -> None: + """Resolve dense and expert GTP degrees against the language grid's TP axes.""" + args.tensor_parallel_num_weight_shards, args.gtp_weight_remat_size = ( + resolve_tensor_parallel_weight_shards( + args.llm_tp, + getattr(args, "tensor_parallel_num_weight_shards", None), + getattr(args, "gtp_weight_remat_size", 1), + ) + ) + ( + args.expert_tensor_parallel_num_weight_shards, + args.expert_gtp_weight_remat_size, + ) = resolve_tensor_parallel_weight_shards( + args.llm_expt_tp or 1, + getattr(args, "expert_tensor_parallel_num_weight_shards", None), + getattr(args, "expert_gtp_weight_remat_size", 1), + shards_field="expert_tensor_parallel_num_weight_shards", + tp_field="expert_tensor_parallel_size", + ) + + def _num_experts(args: argparse.Namespace) -> int: """Resolve MoE expert count from the stock --num-experts arg.""" value = getattr(args, "num_experts", None) diff --git a/examples/mimo/training/data.py b/examples/mimo/training/data.py index 9139711662b..986c521874e 100644 --- a/examples/mimo/training/data.py +++ b/examples/mimo/training/data.py @@ -291,8 +291,12 @@ def build_train_valid_test_data_loaders( raise ValueError(f"unsupported dataset provider: {args.dataset_provider}") encoder_name = _encoder_name(topology) - if encoder_name is not None and (args.micro_batch_size * args.llm_dp) % args.encoder_dp: - raise ValueError("micro_batch_size * llm_dp must be divisible by encoder_dp") + llm_data_parallel_size = args.llm_dp * args.gtp_weight_remat_size + if ( + encoder_name is not None + and (args.micro_batch_size * llm_data_parallel_size) % args.encoder_dp + ): + raise ValueError("micro_batch_size * llm_dp * GTP must be divisible by encoder_dp") language_grid = topology.grids[MIMO_LANGUAGE_MODULE_KEY] language_pgc = topology.module_pgs[MIMO_LANGUAGE_MODULE_KEY] @@ -312,7 +316,7 @@ def build_train_valid_test_data_loaders( if encoder_needs_data and language_needs_data: raise ValueError("the external DataLoader adapter requires non-colocated module grids") if encoder_needs_data: - encoder_mbs = args.micro_batch_size * args.llm_dp // args.encoder_dp + encoder_mbs = args.micro_batch_size * llm_data_parallel_size // args.encoder_dp return _build_split_loaders( args, batch_size=encoder_mbs, @@ -340,7 +344,8 @@ def _build_split_loaders( encoder_name: Optional[str], ) -> tuple[DataLoader, DataLoader, DataLoader]: """Build split-local datasets with deterministic module/DP/split seeds.""" - base_seed = args.seed + module_seed_offset + get_pg_rank(pg_collection.dp) + data_group = pg_collection.dp_cp_gtp_remat or pg_collection.dp + base_seed = args.seed + module_seed_offset + get_pg_rank(data_group) common = _mock_loader_kwargs(args, encoder_name) return tuple( _build_mock_vlm_dataloader( diff --git a/examples/mimo/training/grad_sync.py b/examples/mimo/training/grad_sync.py index 2a06a4b8188..2b04b89c30f 100644 --- a/examples/mimo/training/grad_sync.py +++ b/examples/mimo/training/grad_sync.py @@ -62,7 +62,7 @@ def _is_pg_member(pg) -> bool: def _is_token_source_rank(language_pg) -> bool: """Whether this rank is on the LLM (last PP stage, TP rank 0) coordinate that sums - the global token count over DP/CP. + the global token count over DP/CP/GTP. Sourcing from this single coordinate avoids double-counting across TP/PP replicas. The _is_pg_member guards short-circuit encoder-grid ranks (non-member pp/tp groups) @@ -101,18 +101,18 @@ def _token_source_global_rank(language_grid) -> int: def _global_token_count(num_tokens, language_pg, src_global_rank) -> float: """Total non-padded tokens in the global batch, visible on every rank. - Only the LLM token-source rank computes the count by summing over the LLM DP/CP + Only the LLM token-source ranks compute the count by summing over the LLM DP/CP/GTP group; it then broadcasts that N_global from its global rank to every rank in the world (including the non-colocated encoder grid, where ``language_pg`` is None) so both modules divide by the same per-token mean. """ global_num_tokens = torch.zeros(1, dtype=torch.float32, device="cuda") if _is_token_source_rank(language_pg): - # Collective over DP/CP: every (pp_last, tp0) rank participates so the all-reduce - # does not hang; only DP/CP rank 0 keeps the result and is the broadcast root. + # Collective over every data lane: all (pp_last, tp0) GTP peers participate. + data_group = language_pg.dp_cp_gtp_remat or language_pg.dp_cp token_count = num_tokens.to(dtype=torch.float32).sum().view(1) - dist.all_reduce(token_count, group=language_pg.dp_cp, op=dist.ReduceOp.SUM) - if dist.get_rank(group=language_pg.dp_cp) == 0: + dist.all_reduce(token_count, group=data_group, op=dist.ReduceOp.SUM) + if dist.get_rank(group=data_group) == 0: global_num_tokens.copy_(token_count) dist.broadcast(global_num_tokens, src=src_global_rank) return float(global_num_tokens.item()) diff --git a/examples/mimo/training/runtime.py b/examples/mimo/training/runtime.py index 6a5235aba13..8b46f1fb5f2 100644 --- a/examples/mimo/training/runtime.py +++ b/examples/mimo/training/runtime.py @@ -40,7 +40,7 @@ def configure_module_rng( so disjoint modules (and stages) get independent RNG state. Caller invokes once per active module on this rank. """ - for _required in ("pp", "dp", "tp", "ep", "expt_tp"): + for _required in ("pp", "dp", "tp", "ep", "expt_tp", "gtp_remat", "expt_gtp_remat"): assert ( getattr(pg_collection, _required, None) is not None ), f"pg_collection passed to configure_module_rng must define {_required}" @@ -52,6 +52,8 @@ def configure_module_rng( tp_group=pg_collection.tp, ep_group=pg_collection.ep, etp_group=pg_collection.expt_tp, + gtp_remat_group=pg_collection.gtp_remat, + egtp_remat_group=pg_collection.expt_gtp_remat, ) diff --git a/examples/mimo/training/topology.py b/examples/mimo/training/topology.py index 60e473f58ff..351d6de3dcc 100644 --- a/examples/mimo/training/topology.py +++ b/examples/mimo/training/topology.py @@ -33,23 +33,28 @@ class ModuleGridSpec: cp: int = 1 pp: int = 1 ep: int = 1 + gtp_remat: int = 1 rank_offset: int = 0 # Experts default to TP=1 (set explicitly for MoE); intentionally not Megatron's etp=tp default. expt_tp: int = 1 + expt_gtp_remat: int = 1 dp: int = field(init=False) expt_dp: int = field(init=False) def __post_init__(self) -> None: - dense = self.tp * self.cp * self.pp + dense = self.tp * self.gtp_remat * self.cp * self.pp if self.num_ranks % dense != 0: raise ValueError( - f"num_ranks ({self.num_ranks}) must be divisible by tp*cp*pp ({dense})" + "num_ranks " + f"({self.num_ranks}) must be divisible by tp*gtp_remat*cp*pp ({dense})" ) self.dp = self.num_ranks // dense - expert = self.expt_tp * self.ep * self.pp + expert = self.expt_tp * self.ep * self.expt_gtp_remat * self.pp if self.num_ranks % expert != 0: raise ValueError( - f"num_ranks ({self.num_ranks}) must be divisible by expt_tp*ep*pp ({expert})" + "num_ranks " + f"({self.num_ranks}) must be divisible by " + f"expt_tp*ep*expt_gtp_remat*pp ({expert})" ) self.expt_dp = self.num_ranks // expert @@ -118,27 +123,45 @@ def create_topology(specs: list[ModuleGridSpec]) -> HeteroTopology: def _build_grid(spec: ModuleGridSpec) -> HyperCommGrid: """Create a dense grid plus its expert view and the process groups MIMO needs.""" grid = HyperCommGrid( - shape=[spec.tp, spec.cp, spec.dp, spec.pp], - dim_names=["tp", "cp", "dp", "pp"], + shape=[spec.tp, spec.gtp_remat, spec.cp, spec.dp, spec.pp], + dim_names=["tp", "gtp_remat", "cp", "dp", "pp"], rank_offset=spec.rank_offset, backend="nccl", ) # Expert factorization over the same rank span; pp is shared with the base view. grid.register_view( _EXPERT_VIEW, - shape=[spec.expt_tp, spec.ep, spec.expt_dp, spec.pp], - dim_names=["expt_tp", "ep", "expt_dp", "pp"], + shape=[spec.expt_tp, spec.ep, spec.expt_gtp_remat, spec.expt_dp, spec.pp], + dim_names=["expt_tp", "ep", "expt_gtp_remat", "expt_dp", "pp"], shared_dims=["pp"], ) try: for dims in ( - ["tp"], ["cp"], ["pp"], ["dp"], - ["dp", "cp"], ["tp", "cp"], ["tp", "pp"], - ["tp", "dp"], ["tp", "dp", "cp"], ["tp", "cp", "dp", "pp"], + ["tp"], + ["gtp_remat"], + ["cp"], + ["pp"], + ["dp"], + ["dp", "cp"], + ["gtp_remat", "dp", "cp"], + ["tp", "cp"], + ["tp", "gtp_remat", "pp"], + ["tp", "gtp_remat", "dp"], + ["tp", "gtp_remat", "dp", "cp"], + ["tp", "gtp_remat", "cp", "dp", "pp"], ): grid.create_pg(dims) - for dims in (["ep"], ["expt_tp"], ["expt_dp"], ["expt_tp", "ep"], ["expt_tp", "ep", "pp"]): + for dims in ( + ["ep"], + ["expt_tp"], + ["expt_gtp_remat"], + ["expt_dp"], + ["expt_tp", "ep"], + ["expt_tp", "ep", "pp"], + ["expt_tp", "ep", "expt_gtp_remat", "pp"], + ["expt_gtp_remat", "expt_dp"], + ): grid.create_pg(dims, view=_EXPERT_VIEW) except Exception: grid.destroy() @@ -193,18 +216,27 @@ def pg_collection_from_grid( pgc.pp = grid.get_pg("pp") pgc.dp = grid.get_pg("dp") pgc.dp_cp = grid.get_pg(["dp", "cp"]) + pgc.dp_cp_gtp_remat = grid.get_pg(["gtp_remat", "dp", "cp"]) pgc.intra_dp_cp = pgc.dp_cp + pgc.gtp_remat = grid.get_pg("gtp_remat") pgc.tp_cp = grid.get_pg(["tp", "cp"]) - pgc.tp_dp = grid.get_pg(["tp", "dp"]) - pgc.tp_dp_cp = grid.get_pg(["tp", "dp", "cp"]) - pgc.mp = grid.get_pg(["tp", "pp"]) - pgc.intra_dist_opt = grid.get_pg(["tp", "cp", "dp", "pp"]) + pgc.tp_dp = grid.get_pg(["tp", "gtp_remat", "dp"]) + pgc.tp_dp_cp = grid.get_pg(["tp", "gtp_remat", "dp", "cp"]) + pgc.mp = grid.get_pg(["tp", "gtp_remat", "pp"]) + pgc.intra_dist_opt = grid.get_pg(["tp", "gtp_remat", "cp", "dp", "pp"]) pgc.ep = grid.get_pg("ep", view=_EXPERT_VIEW) pgc.expt_tp = grid.get_pg("expt_tp", view=_EXPERT_VIEW) + pgc.expt_gtp_remat = grid.get_pg("expt_gtp_remat", view=_EXPERT_VIEW) pgc.expt_dp = grid.get_pg("expt_dp", view=_EXPERT_VIEW) + pgc.expt_dp_gtp_remat = grid.get_pg( + ["expt_gtp_remat", "expt_dp"], view=_EXPERT_VIEW + ) pgc.intra_expt_dp = pgc.expt_dp pgc.tp_ep = grid.get_pg(["expt_tp", "ep"], view=_EXPERT_VIEW) pgc.tp_ep_pp = grid.get_pg(["expt_tp", "ep", "pp"], view=_EXPERT_VIEW) + pgc.tp_ep_pp_with_egtp_remat = grid.get_pg( + ["expt_tp", "ep", "expt_gtp_remat", "pp"], view=_EXPERT_VIEW + ) pgc.embd = None pgc.pos_embd = None if is_language: diff --git a/megatron/core/distributed/finalize_model_grads.py b/megatron/core/distributed/finalize_model_grads.py index af81145b6d8..d5be7607714 100644 --- a/megatron/core/distributed/finalize_model_grads.py +++ b/megatron/core/distributed/finalize_model_grads.py @@ -492,7 +492,10 @@ def _allreduce_non_tensor_model_parallel_grads( def _allreduce_replicated_grads_over_gtp_remat_group( - model: List[torch.nn.Module], calculate_per_token_loss: bool = False + model: List[torch.nn.Module], + gtp_remat_group: Optional[torch.distributed.ProcessGroup], + egtp_remat_group: Optional[torch.distributed.ProcessGroup], + calculate_per_token_loss: bool = False, ): """Complete the gtp_remat / egtp_remat axis reduction for replicated parameters. @@ -510,12 +513,6 @@ def _allreduce_replicated_grads_over_gtp_remat_group( No-op when GTP_remat is inactive (group size <= 1). """ - pg_collection = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=["gtp_remat", "expt_gtp_remat"] - ) - gtp_remat_group = pg_collection.gtp_remat - egtp_remat_group = pg_collection.expt_gtp_remat - dense_active = gtp_remat_group is not None and gtp_remat_group.size() > 1 expert_active = egtp_remat_group is not None and egtp_remat_group.size() > 1 if not dense_active and not expert_active: @@ -604,12 +601,29 @@ def finalize_model_grads( # Full DP x CP x gtp_remat group: num_tokens (the per-token-loss divisor below) counts the # gtp_remat peers' distinct tokens. Falls back to replicate dp_cp when gtp is inactive. dp_cp_group = getattr(pg_collection, 'dp_cp_gtp_remat', None) or pg_collection.dp_cp + gtp_remat_group = getattr(pg_collection, 'gtp_remat', None) + egtp_remat_group = getattr(pg_collection, 'expt_gtp_remat', None) else: tp_group = parallel_state.get_tensor_model_parallel_group() pp_group = parallel_state.get_pipeline_model_parallel_group() embd_group = parallel_state.get_embedding_group(check_initialized=False) pos_emb_group = parallel_state.get_position_embedding_group(check_initialized=False) dp_cp_group = parallel_state.get_data_parallel_group(with_context_parallel=True) + gtp_remat_group = parallel_state.get_gtp_weight_remat_group(check_initialized=False) + egtp_remat_group = parallel_state.get_expert_gtp_weight_remat_group(check_initialized=False) + + # A missing group would silently skip the gtp_remat-axis reduction below and train on + # wrong gradients, so fail loudly whenever the config says the axis is active. + for axis, group, axis_size in ( + ('gtp_remat', gtp_remat_group, config.gtp_weight_remat_size), + ('expt_gtp_remat', egtp_remat_group, config.expert_gtp_weight_remat_size), + ): + if axis_size > 1: + found = 'None' if group is None else f'a size-{group.size()} group' + assert group is not None and group.size() == axis_size, ( + f"{axis} is enabled (size={axis_size}) but pg_collection provides {found}. " + f"Pass a pg_collection carrying `{axis}` to finalize_model_grads." + ) # Fence the current stream against all GTP backward grad work before the DP gradient sync. if config.gtp_weight_remat_size > 1 or config.expert_gtp_weight_remat_size > 1: @@ -646,7 +660,10 @@ def finalize_model_grads( ) _allreduce_non_tensor_model_parallel_grads(model, config, tp_group) _allreduce_replicated_grads_over_gtp_remat_group( - model, calculate_per_token_loss=config.calculate_per_token_loss + model, + gtp_remat_group, + egtp_remat_group, + calculate_per_token_loss=config.calculate_per_token_loss, ) if config.timers is not None: config.timers('non-tensor-parallel-grads-all-reduce').stop() diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 1251c85dcee..e72868609bd 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -33,7 +33,7 @@ get_tensor_model_parallel_world_size, model_parallel_is_initialized, ) -from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.process_groups_config import ProcessGroupCollection, resolve_gtp_remat_group from megatron.core.quantization.quant_config import QuantizationConfig from megatron.core.quantization.utils import get_quant_config_or_none from megatron.core.tensor_parallel.layers import ( @@ -1334,6 +1334,8 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): sharded_offsets, tp_group=self._tp_group, dp_cp_group=metadata["dp_cp_group"], + intra_dp_cp_group=metadata.get("intra_dp_cp_group"), + intra_expt_dp_group=metadata.get("intra_expt_dp_group"), ) def backward_dw(self): @@ -1362,10 +1364,13 @@ def __init__( tp_group: Optional[torch.distributed.ProcessGroup] = None, stride: int = 1, name: str | None = None, + pg_collection: Optional[ProcessGroupCollection] = None, ): """ Args: name (str | None): module instance name passed top-down from its paranet module + pg_collection (ProcessGroupCollection | None): process groups used by this layer. + Falls back to the MPU global process groups when not given. """ if not HAVE_TE: raise ImportError( @@ -1457,10 +1462,7 @@ def __init__( ), "Must have at least TE version 2.3 or higher to use symmetric memory all reduce" extra_kwargs["symmetric_ar_type"] = self.config.symmetric_ar_type - pg_collection = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=["gtp_remat", "expt_gtp_remat"] - ) - gtp_remat_group = pg_collection.expt_gtp_remat if is_expert else pg_collection.gtp_remat + gtp_remat_group = resolve_gtp_remat_group(pg_collection, is_expert) self.stride = stride self.te_quant_params: Optional[TEQuantizationParams] = None @@ -1578,6 +1580,8 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): sharded_offsets, tp_group=self._tp_group, dp_cp_group=metadata["dp_cp_group"], + intra_dp_cp_group=metadata.get("intra_dp_cp_group"), + intra_expt_dp_group=metadata.get("intra_expt_dp_group"), ) @override @@ -1621,10 +1625,13 @@ def __init__( tp_group: Optional[torch.distributed.ProcessGroup] = None, stride: int = 1, name: str | None = None, + pg_collection: Optional[ProcessGroupCollection] = None, ): """ Args: name (str | None): module instance name passed top-down from its paranet module + pg_collection (ProcessGroupCollection | None): process groups used by this layer. + Falls back to the MPU global process groups when not given. """ if not HAVE_TE: raise ImportError( @@ -1639,10 +1646,7 @@ def __init__( world_size = get_pg_size(tp_group) rank = get_pg_rank(tp_group) self.stride = stride - pg_collection = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=["gtp_remat", "expt_gtp_remat"] - ) - gtp_remat_group = pg_collection.expt_gtp_remat if is_expert else pg_collection.gtp_remat + gtp_remat_group = resolve_gtp_remat_group(pg_collection, is_expert) super().__init__( input_size=input_size, @@ -1712,6 +1716,8 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): sharded_offsets, tp_group=self._tp_group, dp_cp_group=metadata["dp_cp_group"], + intra_dp_cp_group=metadata.get("intra_dp_cp_group"), + intra_expt_dp_group=metadata.get("intra_expt_dp_group"), ) @override @@ -1882,10 +1888,13 @@ def __init__( tp_comm_buffer_name: Optional[str] = None, tp_group: Optional[torch.distributed.ProcessGroup] = None, name: str | None = None, + pg_collection: Optional[ProcessGroupCollection] = None, ): """ Args: name (str | None): module instance name passed top-down from its paranet module + pg_collection (ProcessGroupCollection | None): process groups used by this layer. + Falls back to the MPU global process groups when not given. """ if not HAVE_TE: raise ImportError( @@ -1899,10 +1908,7 @@ def __init__( ) tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert) self._tp_group = tp_group - pg_collection = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=["gtp_remat", "expt_gtp_remat"] - ) - gtp_remat_group = pg_collection.expt_gtp_remat if is_expert else pg_collection.gtp_remat + gtp_remat_group = resolve_gtp_remat_group(pg_collection, is_expert) super().__init__( input_size=input_size, @@ -1969,6 +1975,8 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): sharded_offsets, tp_group=self._tp_group, dp_cp_group=metadata["dp_cp_group"], + intra_dp_cp_group=metadata.get("intra_dp_cp_group"), + intra_expt_dp_group=metadata.get("intra_expt_dp_group"), ) @override @@ -2323,6 +2331,8 @@ def sharded_state_dict( sharded_offsets, tp_group=self._tp_group, dp_cp_group=metadata["dp_cp_group"], + intra_dp_cp_group=metadata.get("intra_dp_cp_group"), + intra_expt_dp_group=metadata.get("intra_expt_dp_group"), ) @@ -2783,6 +2793,8 @@ def get_gemm_tensor(param_name: str, gemm_idx: int) -> torch.Tensor: new_sharded_offsets, tp_group=self._tp_group, dp_cp_group=metadata["dp_cp_group"], + intra_dp_cp_group=metadata.get("intra_dp_cp_group"), + intra_expt_dp_group=metadata.get("intra_expt_dp_group"), ) # Remove expert layers indexing from sharded keys replace_prefix_for_sharding(sub_sd, f"{gemm_idx}.", expert_prefix) diff --git a/megatron/core/models/common/embeddings/language_model_embedding.py b/megatron/core/models/common/embeddings/language_model_embedding.py index 7e49ec6c02d..08f5d359082 100644 --- a/megatron/core/models/common/embeddings/language_model_embedding.py +++ b/megatron/core/models/common/embeddings/language_model_embedding.py @@ -6,6 +6,7 @@ from torch import Tensor from megatron.core import tensor_parallel +from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import get_tensor_model_parallel_group_if_none, nvtx_decorator @@ -35,6 +36,7 @@ def __init__( num_tokentypes: int = 0, scatter_to_sequence_parallel: bool = True, tp_group: Optional[torch.distributed.ProcessGroup] = None, + pg_collection: Optional[ProcessGroupCollection] = None, ): super().__init__(config=config) @@ -60,6 +62,7 @@ def __init__( reduce_scatter_embeddings=self.reduce_scatter_embeddings, config=self.config, tp_group=self.tp_group, + pg_collection=pg_collection, ) # Position embedding (serial). diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py index 53522dd8b2b..51ab606dd5a 100644 --- a/megatron/core/models/common/language_module/language_module.py +++ b/megatron/core/models/common/language_module/language_module.py @@ -519,4 +519,6 @@ def tie_embeddings_and_output_weights_state_dict( allow_shape_mismatch=True, tp_group=self.tp_group, dp_cp_group=metadata['dp_cp_group'], + intra_dp_cp_group=metadata.get('intra_dp_cp_group'), + intra_expt_dp_group=metadata.get('intra_expt_dp_group'), ) diff --git a/megatron/core/models/hybrid/hybrid_model.py b/megatron/core/models/hybrid/hybrid_model.py index f750c77e05b..f0358de57b9 100644 --- a/megatron/core/models/hybrid/hybrid_model.py +++ b/megatron/core/models/hybrid/hybrid_model.py @@ -239,6 +239,7 @@ def __init__( position_embedding_type=position_embedding_type, scatter_to_sequence_parallel=scatter_embedding_sequence_parallel, tp_group=self.pg_collection.tp, + pg_collection=self.pg_collection, ) # MLA (also used by DeepSeek Sparse Attention) uses its own decoupled RoPE, therefore we do @@ -322,6 +323,7 @@ def __init__( skip_weight_param_allocation=self.pre_process and self.share_embeddings_and_output_weights, tp_group=self.pg_collection.tp, + pg_collection=self.pg_collection, ) if self.pre_process or self.post_process or self.mtp_process: diff --git a/megatron/core/models/mimo/model/base.py b/megatron/core/models/mimo/model/base.py index b226b5c1e4b..5f5c25fa281 100644 --- a/megatron/core/models/mimo/model/base.py +++ b/megatron/core/models/mimo/model/base.py @@ -125,11 +125,14 @@ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None): pg = getattr(pg_src, 'pg_collection', None) mod_metadata = metadata if pg is not None: - assert ( - hasattr(pg, 'dp_cp') and pg.dp_cp is not None - ), f"pg_collection on '{name}' is missing dp_cp group" + dp_cp_group = pg.dp_cp_gtp_remat or pg.dp_cp + assert dp_cp_group is not None, ( + f"pg_collection on '{name}' is missing a data-parallel group" + ) mod_metadata = dict(metadata) if metadata else {} - mod_metadata['dp_cp_group'] = pg.dp_cp + mod_metadata['dp_cp_group'] = dp_cp_group + mod_metadata['intra_dp_cp_group'] = pg.dp_cp + mod_metadata['intra_expt_dp_group'] = pg.expt_dp # Unwrap wrappers so the sharded keys match the raw load_state_dict keys. inner = module child_prefix = f'{prefix}{name}.' @@ -549,6 +552,10 @@ def _attach_modality_split_sizes( language_grid = grid_map[MIMO_LANGUAGE_MODULE_KEY] encoder_dp = encoder_grid.shape[encoder_grid.dim_names.index("dp")] language_dp = language_grid.shape[language_grid.dim_names.index("dp")] + if "gtp_remat" in language_grid.dim_names: + language_dp *= language_grid.shape[ + language_grid.dim_names.index("gtp_remat") + ] assert encoder_dp <= language_dp, ( f"Bridge fan-out split metadata with non-uniform per-sample sizes " f"requires encoder DP <= LM DP (got encoder='{encoder_name}' " diff --git a/megatron/core/models/mimo/optimizer.py b/megatron/core/models/mimo/optimizer.py index 598f6c883af..7d488d35487 100644 --- a/megatron/core/models/mimo/optimizer.py +++ b/megatron/core/models/mimo/optimizer.py @@ -16,6 +16,7 @@ from megatron.core.optimizer.optimizer import MegatronOptimizer from megatron.core.optimizer.optimizer_config import OptimizerConfig from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.utils import unwrap_model if TYPE_CHECKING: from megatron.core.hyper_comm_grid import HyperCommGrid @@ -322,7 +323,7 @@ def _restore_grad_scaler(sub_sd): def _get_replica_id(pg_collection: Optional[ProcessGroupCollection]) -> tuple: """Build replica_id tuple for ShardedObject deduplication. - Returns (tp_rank, pp_rank, dp_rank) so only (0, 0, 0) within each + Returns (tp_gtp_rank, pp_rank, dp_rank) so only (0, 0, 0) within each module's parallelism group is the main replica; all other ranks in the same module are non-main replicas of the same object. """ @@ -336,30 +337,10 @@ def _get_replica_id(pg_collection: Optional[ProcessGroupCollection]) -> tuple: assert ( hasattr(pg_collection, 'dp') and pg_collection.dp is not None ), "pg_collection.dp must be set for checkpoint deduplication" - return (pg_collection.tp.rank(), pg_collection.pp.rank(), pg_collection.dp.rank()) - - -_EXPERT_VIEW = "expert" - - -def _get_pg_collection_for_optimizer(grid) -> ProcessGroupCollection: - """Derive the optimizer's ProcessGroupCollection from a populated HyperCommGrid. - - Dense groups come from the base view; expert-parallel groups (tp_ep_pp, expt_dp) come from - the grid's dedicated expert view -- expert parallelism is always factored into a separate - view (expt_tp/ep/expt_dp), never the base view. All groups must be pre-created on the grid. - """ - pg = ProcessGroupCollection() - pg.dp = grid.get_pg("dp") - pg.dp_cp = grid.get_pg(["dp", "cp"]) - pg.tp = grid.get_pg("tp") - pg.pp = grid.get_pg("pp") - pg.mp = grid.get_pg(["tp", "pp"]) - pg.tp_ep_pp = grid.get_pg(["expt_tp", "ep", "pp"], view=_EXPERT_VIEW) - pg.expt_dp = grid.get_pg("expt_dp", view=_EXPERT_VIEW) - # Distributed-optimizer grad-stats group spans the dense shards (mirrors the topology PGC). - pg.intra_dist_opt = grid.get_pg(["tp", "cp", "dp", "pp"]) - return pg + gtp_group = pg_collection.gtp_remat + assert gtp_group is not None, "pg_collection.gtp_remat must be set for checkpoint deduplication" + tp_gtp_rank = pg_collection.tp.rank() * gtp_group.size() + gtp_group.rank() + return (tp_gtp_rank, pg_collection.pp.rank(), pg_collection.dp.rank()) def get_mimo_optimizer(mimo_model: "MimoModel", config: OptimizerConfig) -> MimoOptimizer: @@ -386,7 +367,10 @@ def get_mimo_optimizer(mimo_model: "MimoModel", config: OptimizerConfig) -> Mimo module = mimo_model.modality_submodules[module_name] if module is not None: - pg_collection = _get_pg_collection_for_optimizer(grid) + pg_collection = getattr(unwrap_model(module), 'pg_collection', None) + assert pg_collection is not None, ( + f"Module '{module_name}' must own a ProcessGroupCollection before optimizer setup" + ) assert ( not hasattr(module, 'ddp_config') or module.ddp_config is None diff --git a/megatron/core/models/mimo/submodules/base.py b/megatron/core/models/mimo/submodules/base.py index ac7bf64c063..1e0c413d9f1 100644 --- a/megatron/core/models/mimo/submodules/base.py +++ b/megatron/core/models/mimo/submodules/base.py @@ -82,11 +82,12 @@ def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None): parallel_state global fallback in ensure_metadata_has_dp_cp_group. """ if self.pg_collection is not None: - assert ( - hasattr(self.pg_collection, 'dp_cp') and self.pg_collection.dp_cp is not None - ), "pg_collection is missing dp_cp group" + dp_cp_group = self.pg_collection.dp_cp_gtp_remat or self.pg_collection.dp_cp + assert dp_cp_group is not None, "pg_collection is missing a data-parallel group" metadata = dict(metadata) if metadata else {} - metadata['dp_cp_group'] = self.pg_collection.dp_cp + metadata['dp_cp_group'] = dp_cp_group + metadata['intra_dp_cp_group'] = self.pg_collection.dp_cp + metadata['intra_expt_dp_group'] = self.pg_collection.expt_dp sharded_sd = {} for name, container in self.named_children(): diff --git a/megatron/core/pipeline_parallel/bridge_communicator.py b/megatron/core/pipeline_parallel/bridge_communicator.py index fc234cac8ae..8b806214d39 100644 --- a/megatron/core/pipeline_parallel/bridge_communicator.py +++ b/megatron/core/pipeline_parallel/bridge_communicator.py @@ -14,11 +14,11 @@ class CommRole(Enum): """Communication role for ranks in bridge communication. - SENDER: Leader tp-cp rank within each DP replica of source grid. + SENDER: Leader tp-cp rank within each data lane of source grid. Sends data to destination grid receivers. - RECEIVER: Leader tp-cp rank within each DP replica of destination grid. + RECEIVER: Leader tp-cp rank within each data lane of destination grid. Receives data from source grid senders. - MEMBER: Non-leader ranks within DP replicas. + MEMBER: Non-leader ranks within data lanes. Participate in broadcasts from their local leader. """ @@ -212,44 +212,48 @@ def _get_or_create_bridge_pg(cls, ranks: List[int]): def get_leader_rank(self, grid: HyperCommGrid, is_src: bool) -> List[int]: """Get the leader rank for a given grid and direction. - We elect leader rank for each dp replica, the first tp-cp rank in the group + We elect one leader for each DP x GTP data lane, the first tp-cp rank in the group in the last pp stage (for src grid) or first pp stage (for dest grid) is the leader. """ leader_ranks = [] local_leader_rank = None - # grid.gen_rank_enum(["tp", "cp", "pp"]) # vary tp & cp, but same dp + # Vary model dimensions while holding data-like dimensions fixed. GTP ranks execute + # independent microbatches after rematerializing the same TP weight slice. + data_dims = {"dp", "gtp_remat"} # returns a list of sublists, each sublist is a group of ranks - # that have different tp & cp & pp, same dp - per_dp_replica_ranks = grid._gen_rank_enum([x for x in grid.dim_names if x != "dp"]) + # that have different tp & cp & pp, same dp and gtp_remat + per_data_lane_ranks = grid._gen_rank_enum( + [dim for dim in grid.dim_names if dim not in data_dims] + ) if is_src: # Add rank from last pp stage ranks = [] - for group in per_dp_replica_ranks: + for group in per_data_lane_ranks: if self.current_rank in group: assert ( local_leader_rank is None - ), "only one local leader rank is allowed per dp replica" + ), "only one local leader rank is allowed per data lane" local_leader_rank = group[-1] ranks.append(group[-1]) leader_ranks.extend(ranks) else: # Add rank from first pp stage ranks = [] - for group in per_dp_replica_ranks: + for group in per_data_lane_ranks: if self.current_rank in group: assert ( local_leader_rank is None - ), "only one local leader rank is allowed per dp replica" + ), "only one local leader rank is allowed per data lane" local_leader_rank = group[0] ranks.append(group[0]) leader_ranks.extend(ranks) return leader_ranks, local_leader_rank def get_boundary_pp_stage_ranks(self, grid: HyperCommGrid, is_src: bool): - """Get TP-CP ranks at boundary PP stage for each DP replica. + """Get TP-CP ranks at the boundary PP stage for each data lane. Returns ranks at the last PP stage (if src) or first PP stage (if dest) - for each DP dimension, ordered by DP dimension. + with the DP and GTP coordinates held fixed. """ # Get tp-cp rank enumeration (each list has same dp and pp, different tp and cp) diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py index 29d36ab2c7d..1daeacc9027 100644 --- a/megatron/core/pipeline_parallel/schedules.py +++ b/megatron/core/pipeline_parallel/schedules.py @@ -690,6 +690,15 @@ def _build_default_pg_collection() -> ProcessGroupCollection: pg_collection.dp = parallel_state.get_data_parallel_group( with_context_parallel=False, partial_data_parallel=False ) + # gtp_remat axis: consumers read these with getattr and silently skip the gtp_remat + # reduction when absent, so populate them even when GTP_remat is inactive. + pg_collection.gtp_remat = parallel_state.get_gtp_weight_remat_group(check_initialized=False) + pg_collection.expt_gtp_remat = parallel_state.get_expert_gtp_weight_remat_group( + check_initialized=False + ) + pg_collection.dp_cp_gtp_remat = parallel_state.get_data_parallel_group( + with_context_parallel=True, partial_data_parallel=False + ) return pg_collection @@ -1614,7 +1623,7 @@ def forward_backward_helper_wrapper( recv_next = True if is_pp_last_stage(p2p_communicator.pp_group): recv_next = False - (input_tensor, output_tensor_grad) = ( + input_tensor, output_tensor_grad = ( p2p_communicator.send_forward_backward_recv_forward_backward( output_tensor, input_tensor_grad, @@ -1678,7 +1687,7 @@ def forward_backward_helper_wrapper( if is_pp_last_stage(p2p_communicator.pp_group): recv_next = False - (bwd_recv_buffer[-1], bwd_wait_handles) = ( + bwd_recv_buffer[-1], bwd_wait_handles = ( p2p_communicator.send_backward_recv_backward( input_tensor_grad, recv_next=recv_next, @@ -1831,7 +1840,7 @@ def pp_post_backward(input_tensor_grad, vp_stage=None): backward_k, forward=False ) - (bwd_recv_buffer[backward_k % bwd_recv_buffer_size], bwd_wait_handles) = ( + bwd_recv_buffer[backward_k % bwd_recv_buffer_size], bwd_wait_handles = ( p2p_communicator.send_backward_recv_backward( input_tensor_grad, recv_next=recv_next, @@ -1904,7 +1913,7 @@ def pp_post_backward(input_tensor_grad, vp_stage=None): recv_prev = False # Communicate tensors. - (input_tensor, output_tensor_grad) = ( + input_tensor, output_tensor_grad = ( p2p_communicator.send_forward_backward_recv_forward_backward( output_tensor, input_tensor_grad, diff --git a/megatron/core/process_groups_config.py b/megatron/core/process_groups_config.py index ccb6dce0eb8..970713941ae 100644 --- a/megatron/core/process_groups_config.py +++ b/megatron/core/process_groups_config.py @@ -689,6 +689,31 @@ def setup_process_groups_for_ddp( return result +def resolve_gtp_remat_group( + pg_collection: Optional["ProcessGroupCollection"], is_expert: bool +) -> Optional[torch.distributed.ProcessGroup]: + """Resolve the gtp_remat / expt_gtp_remat group for a weight-owning module. + + Prefers the group carried by ``pg_collection``; falls back to the MPU globals when the + caller passed no collection, or one predating the gtp_remat fields. The fallback keeps + pre-pg_collection callers working — a collection that does carry the field is always + honored, including when it holds a custom (non-MPU) group. + + Args: + pg_collection: Collection supplied by the caller, or None. + is_expert: Select the expert axis (``expt_gtp_remat``) instead of the dense one. + """ + attr = 'expt_gtp_remat' if is_expert else 'gtp_remat' + # `vars()`, not hasattr: __getattr__ makes hasattr always True, so the fallback below + # would be unreachable. + if pg_collection is not None and attr in vars(pg_collection): + return getattr(pg_collection, attr) + mpu_pgs = ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['gtp_remat', 'expt_gtp_remat'] + ) + return getattr(mpu_pgs, attr) + + @dataclass class MultiModuleProcessGroupCollection: """Process group collection for multi-module pipelines. diff --git a/megatron/core/ssm/mamba_mixer.py b/megatron/core/ssm/mamba_mixer.py index f6ae07dd230..3d2158818a0 100644 --- a/megatron/core/ssm/mamba_mixer.py +++ b/megatron/core/ssm/mamba_mixer.py @@ -128,6 +128,8 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): sharded_offsets, tp_group=self.tp_group, dp_cp_group=metadata["dp_cp_group"], + intra_dp_cp_group=metadata.get("intra_dp_cp_group"), + intra_expt_dp_group=metadata.get("intra_expt_dp_group"), ) @@ -291,6 +293,7 @@ def __init__( is_expert=False, tp_comm_buffer_name="fc1", tp_group=self.pg_collection.tp, + pg_collection=self.pg_collection, name=(name + f".in_proj") if name is not None else None, ) # in_proj packs [z, x, B, C, dt] into one ColumnParallelLinear. Each @@ -442,6 +445,7 @@ def __init__( is_expert=False, tp_comm_buffer_name="fc2", tp_group=self.pg_collection.tp, + pg_collection=self.pg_collection, name=(name + f".out_proj") if name is not None else None, ) @@ -1430,6 +1434,8 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): sharded_offsets=sharded_offsets, tp_group=self.tp_group, dp_cp_group=metadata["dp_cp_group"], + intra_dp_cp_group=metadata.get("intra_dp_cp_group"), + intra_expt_dp_group=metadata.get("intra_expt_dp_group"), ) # Submodules for name, module in self.named_children(): @@ -1481,6 +1487,8 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): prepend_offsets=sharded_offsets, tp_group=self.tp_group, dp_cp_group=metadata['dp_cp_group'], + intra_dp_cp_group=metadata.get('intra_dp_cp_group'), + intra_expt_dp_group=metadata.get('intra_expt_dp_group'), ) assert sharded_state_dict[f"{prefix}in_proj.weight"].data.size(0) == in_proj_dim, ( diff --git a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py b/megatron/core/tensor_parallel/generalized_tensor_parallelism.py index 9099cd114b9..a5db7cc1093 100644 --- a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py +++ b/megatron/core/tensor_parallel/generalized_tensor_parallelism.py @@ -2149,6 +2149,7 @@ def make_sharded_tensors_for_checkpoint_with_gtp_remat( tp_group, dp_cp_group, intra_dp_cp_group=None, + intra_expt_dp_group=None, ): """GTP-aware analogue of make_sharded_tensors_for_checkpoint. @@ -2190,16 +2191,22 @@ def make_sharded_tensors_for_checkpoint_with_gtp_remat( gtp_rank = get_pg_rank(gtp_remat_group) gtp_remat_size = get_pg_size(gtp_remat_group) - # Replicate-group rank — the true replicas of a given GTP chunk live here. - if intra_dp_cp_group is not None: - dp_replica_rank = get_pg_rank(intra_dp_cp_group) - else: + def replica_rank(tensor): + """Rank in the GTP-excluded dense or expert replica group.""" + replica_group = intra_dp_cp_group + if not getattr(tensor, 'allreduce', True) and intra_expt_dp_group is not None: + replica_group = intra_expt_dp_group + if replica_group is not None: + return get_pg_rank(replica_group) from megatron.core import parallel_state # noqa: E402 - dp_replica_rank = parallel_state.get_data_parallel_rank( + return parallel_state.get_data_parallel_rank( with_context_parallel=True, with_gtp_remat=False ) + # Extra-state objects carry no tensor attributes; use the module's GTP weight kind. + gtp_template = next(t for t in state_dict.values() if is_gtp_param(t)) + sharded_state_dict = {} for layer_name, tensor in state_dict.items(): layer_key = f"{prefix}{layer_name}" @@ -2209,7 +2216,11 @@ def make_sharded_tensors_for_checkpoint_with_gtp_remat( # ShardedObject (extra_state metadata): GTP-REPLICATED across the GTP group. Fold # gtp_rank into position 1 of the replica_id (PP, TP-replica-coord, DP) tuple so # GTP-peer ranks within the same TP slice get unique replica_ids. - replica_id = (0, tp_rank * gtp_remat_size + gtp_rank, dp_replica_rank) + replica_id = ( + 0, + tp_rank * gtp_remat_size + gtp_rank, + replica_rank(gtp_template), + ) sharded_state_dict[layer_key] = make_sharded_object_for_checkpoint( tensor, layer_key, sharded_offsets, replica_id=replica_id ) @@ -2220,7 +2231,7 @@ def make_sharded_tensors_for_checkpoint_with_gtp_remat( # ranks would collide on the same replica_id. Inject gtp_rank into replica_id # position 1 (same as the GTP-sharded branch below). if layer_name in tensor_parallel_layers_axis_map: - replica_id = (0, gtp_rank, dp_replica_rank) + replica_id = (0, gtp_rank, replica_rank(tensor)) sharded_state_dict[layer_key] = make_tp_sharded_tensor_for_checkpoint( tensor, layer_key, @@ -2229,9 +2240,15 @@ def make_sharded_tensors_for_checkpoint_with_gtp_remat( prepend_offsets=sharded_offsets, tp_group=tp_group, dp_cp_group=dp_cp_group, + intra_dp_cp_group=intra_dp_cp_group, + intra_expt_dp_group=intra_expt_dp_group, ) else: - replica_id = (0, tp_rank * gtp_remat_size + gtp_rank, dp_replica_rank) + replica_id = ( + 0, + tp_rank * gtp_remat_size + gtp_rank, + replica_rank(tensor), + ) sharded_state_dict[layer_key] = make_sharded_tensor_for_checkpoint( tensor, layer_key, @@ -2253,6 +2270,8 @@ def make_sharded_tensors_for_checkpoint_with_gtp_remat( prepend_offsets=sharded_offsets, tp_group=tp_group, dp_cp_group=dp_cp_group, + intra_dp_cp_group=intra_dp_cp_group, + intra_expt_dp_group=intra_expt_dp_group, ) return sharded_state_dict diff --git a/megatron/core/tensor_parallel/inference_layers.py b/megatron/core/tensor_parallel/inference_layers.py index 87ba3023d2a..1bf0ea8e74a 100644 --- a/megatron/core/tensor_parallel/inference_layers.py +++ b/megatron/core/tensor_parallel/inference_layers.py @@ -20,6 +20,7 @@ from megatron.core.inference.quantization.utils import mm_mxfp8 from megatron.core.inference.symmetric_memory import SymmetricMemoryManager from megatron.core.model_parallel_config import ModelParallelConfig +from megatron.core.process_groups_config import ProcessGroupCollection, resolve_gtp_remat_group from megatron.core.tensor_parallel.mappings import ( gather_from_tensor_model_parallel_region, reduce_scatter_to_sequence_parallel_region, @@ -91,6 +92,7 @@ def __init__( symmetric_ar_type: Optional[str] = None, tp_group: Optional[torch.distributed.ProcessGroup] = None, name: str | None = None, + pg_collection: Optional[ProcessGroupCollection] = None, ): assert HAVE_TE, "--transformer-impl=inference_optimized requires transformer engine" super().__init__( @@ -107,6 +109,8 @@ def __init__( symmetric_ar_type=symmetric_ar_type, tp_group=tp_group, name=name, + # TELinear takes the resolved group rather than the collection. + gtp_remat_group=resolve_gtp_remat_group(pg_collection, is_expert), ) def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, None]: @@ -139,6 +143,7 @@ def __init__( tp_comm_buffer_name: Optional[str] = None, tp_group: Optional[torch.distributed.ProcessGroup] = None, name: str | None = None, + pg_collection: Optional[ProcessGroupCollection] = None, ): assert HAVE_TE, "--transformer-impl=inference_optimized requires transformer engine" super().__init__( @@ -155,6 +160,7 @@ def __init__( tp_comm_buffer_name=tp_comm_buffer_name, tp_group=tp_group, name=name, + pg_collection=pg_collection, ) self.tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert) self.tp_size = dist.get_world_size(self.tp_group) @@ -268,6 +274,7 @@ def __init__( tp_comm_buffer_name: Optional[str] = None, tp_group: Optional[torch.distributed.ProcessGroup] = None, name: str | None = None, + pg_collection: Optional[ProcessGroupCollection] = None, ): assert HAVE_TE, "--transformer-impl=inference_optimized requires transformer engine" super().__init__( @@ -284,6 +291,7 @@ def __init__( tp_comm_buffer_name=tp_comm_buffer_name, tp_group=tp_group, name=name, + pg_collection=pg_collection, ) self.tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert) self.tp_size = dist.get_world_size(self.tp_group) @@ -366,6 +374,7 @@ def __init__( tp_comm_buffer_name: Optional[str] = None, tp_group: Optional[torch.distributed.ProcessGroup] = None, name: str | None = None, + pg_collection: Optional[ProcessGroupCollection] = None, ): assert HAVE_TE, "--transformer-impl=inference_optimized requires transformer engine" super().__init__( @@ -380,6 +389,7 @@ def __init__( tp_comm_buffer_name=tp_comm_buffer_name, tp_group=tp_group, name=name, + pg_collection=pg_collection, ) self.tp_group = get_tensor_model_parallel_group_if_none(tp_group, is_expert=is_expert) self.tp_size = dist.get_world_size(self.tp_group) diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py index a42cbd05841..3f6ad1b3320 100644 --- a/megatron/core/tensor_parallel/layers.py +++ b/megatron/core/tensor_parallel/layers.py @@ -22,7 +22,7 @@ get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) -from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.process_groups_config import ProcessGroupCollection, resolve_gtp_remat_group from megatron.core.utils import ( divide, get_pg_rank, @@ -251,6 +251,7 @@ def __init__( reduce_scatter_embeddings: bool = False, config: ModelParallelConfig, tp_group: Optional[torch.distributed.ProcessGroup] = None, + pg_collection: Optional[ProcessGroupCollection] = None, ): super(VocabParallelEmbedding, self).__init__() # Keep the input dimensions. @@ -261,7 +262,7 @@ def __init__( self.tp_group = get_tensor_model_parallel_group_if_none(self.tp_group) - (self.vocab_start_index, self.vocab_end_index) = ( + self.vocab_start_index, self.vocab_end_index = ( VocabUtility.vocab_range_from_global_vocab_size( self.num_embeddings, get_pg_rank(self.tp_group), get_pg_size(self.tp_group) ) @@ -314,9 +315,7 @@ def __init__( ) self.gtp_remat_size = 1 - gtp_remat_group = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=["gtp_remat"] - ).gtp_remat + gtp_remat_group = resolve_gtp_remat_group(pg_collection, is_expert=False) if gtp_remat_group is not None and gtp_remat_group.size() > 1: from megatron.core.tensor_parallel.gtp_api import wrap_module_params_gtp @@ -397,6 +396,8 @@ def sharded_state_dict( prepend_offsets=sharded_offsets, tp_group=self.tp_group, dp_cp_group=metadata["dp_cp_group"], + intra_dp_cp_group=metadata.get("intra_dp_cp_group"), + intra_expt_dp_group=metadata.get("intra_expt_dp_group"), ) } @@ -935,6 +936,7 @@ def __init__( disable_grad_reduce: bool = False, tp_group: Optional[torch.distributed.ProcessGroup] = None, name: str | None = None, + pg_collection: Optional[ProcessGroupCollection] = None, ): super(ColumnParallelLinear, self).__init__() @@ -1020,10 +1022,7 @@ def __init__( self.weight = None self.gtp_remat_size = 1 - _pg = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=["gtp_remat", "expt_gtp_remat"] - ) - gtp_remat_group = _pg.expt_gtp_remat if self.is_expert else _pg.gtp_remat + gtp_remat_group = resolve_gtp_remat_group(pg_collection, self.is_expert) if gtp_remat_group is not None and gtp_remat_group.size() > 1: from megatron.core.tensor_parallel.gtp_api import wrap_module_params_gtp @@ -1225,6 +1224,8 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): sharded_offsets, tp_group=self.tp_group, dp_cp_group=metadata['dp_cp_group'], + intra_dp_cp_group=metadata.get('intra_dp_cp_group'), + intra_expt_dp_group=metadata.get('intra_expt_dp_group'), ) def set_extra_state(self, state: Any): @@ -1299,6 +1300,7 @@ def __init__( tp_comm_buffer_name: str | None = None, # Not used tp_group: Optional[torch.distributed.ProcessGroup] = None, name: str | None = None, + pg_collection: Optional[ProcessGroupCollection] = None, ): super(RowParallelLinear, self).__init__() @@ -1385,10 +1387,7 @@ def __init__( setattr(self.weight, "allreduce", not use_expert_pgs) self.gtp_remat_size = 1 - _pg = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=["gtp_remat", "expt_gtp_remat"] - ) - gtp_remat_group = _pg.expt_gtp_remat if self.is_expert else _pg.gtp_remat + gtp_remat_group = resolve_gtp_remat_group(pg_collection, self.is_expert) if gtp_remat_group is not None and gtp_remat_group.size() > 1: from megatron.core.tensor_parallel.gtp_api import wrap_module_params_gtp @@ -1505,6 +1504,8 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): sharded_offsets, tp_group=self.tp_group, dp_cp_group=metadata['dp_cp_group'], + intra_dp_cp_group=metadata.get('intra_dp_cp_group'), + intra_expt_dp_group=metadata.get('intra_expt_dp_group'), ) def set_extra_state(self, state: Any): diff --git a/megatron/core/tensor_parallel/random.py b/megatron/core/tensor_parallel/random.py index eb726c2eaf4..22aaf515a8f 100644 --- a/megatron/core/tensor_parallel/random.py +++ b/megatron/core/tensor_parallel/random.py @@ -451,7 +451,11 @@ def model_parallel_cuda_manual_seed( tp_rank: Optional[int] = None, ep_rank: Optional[int] = None, etp_rank: Optional[int] = None, + gtp_remat_rank: Optional[int] = None, + egtp_remat_rank: Optional[int] = None, force_reset_rng: bool = False, + gtp_remat_world_size: Optional[int] = None, + egtp_remat_world_size: Optional[int] = None, ): """Initialize model parallel cuda seed. @@ -476,6 +480,14 @@ def model_parallel_cuda_manual_seed( ep_rank = get_expert_model_parallel_rank() if etp_rank is None: etp_rank = get_expert_tensor_parallel_rank() + if gtp_remat_rank is None: + gtp_remat_rank = get_gtp_weight_remat_rank() + if egtp_remat_rank is None: + egtp_remat_rank = get_expert_gtp_weight_remat_rank() + if gtp_remat_world_size is None: + gtp_remat_world_size = get_gtp_weight_remat_world_size() + if egtp_remat_world_size is None: + egtp_remat_world_size = get_expert_gtp_weight_remat_world_size() # 2718 is just for fun and any POSITIVE value will work. offset = seed + 2718 tensor_model_parallel_seed = offset + tp_rank @@ -500,12 +512,10 @@ def model_parallel_cuda_manual_seed( # must draw DIFFERENT values (everything above is identical across peers by design). The 65536 # stride keeps these disjoint from the tp/ep/etp seeds. Added only when the axis is active, so # non-GTP runs keep a byte-identical tracker set (and checkpoint rng payload). - gtp_remat_rank = get_gtp_weight_remat_rank() - if get_gtp_weight_remat_world_size() > 1: + if gtp_remat_world_size > 1: gtp_remat_seed = tensor_model_parallel_seed + 65536 * (1 + gtp_remat_rank) _CUDA_RNG_STATE_TRACKER.add(_GTP_REMAT_RNG_TRACKER_NAME, gtp_remat_seed) - egtp_remat_rank = get_expert_gtp_weight_remat_rank() - if get_expert_gtp_weight_remat_world_size() > 1: + if egtp_remat_world_size > 1: egtp_remat_seed = expert_parallel_seed + 32768 + 65536 * (1 + egtp_remat_rank) _CUDA_RNG_STATE_TRACKER.add(_EXPERT_GTP_REMAT_RNG_TRACKER_NAME, egtp_remat_seed) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 4ce5babb9c0..682b75fb701 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -421,6 +421,7 @@ def __init__( is_expert=False, tp_comm_buffer_name='proj', tp_group=self.pg_collection.tp, + pg_collection=self.pg_collection, name=(name + ".linear_proj") if name is not None else None, ) @@ -1689,6 +1690,7 @@ def __init__( is_expert=False, tp_comm_buffer_name='qkv', tp_group=self.pg_collection.tp, + pg_collection=self.pg_collection, name=(name + ".linear_qkv") if name is not None else None, ) diff --git a/megatron/core/transformer/dot_product_attention.py b/megatron/core/transformer/dot_product_attention.py index 69039e0bfd0..0e29bbcfa14 100644 --- a/megatron/core/transformer/dot_product_attention.py +++ b/megatron/core/transformer/dot_product_attention.py @@ -268,4 +268,6 @@ def sharded_state_dict( sharded_offsets, tp_group=self.tp_group, dp_cp_group=metadata['dp_cp_group'], + intra_dp_cp_group=metadata.get('intra_dp_cp_group'), + intra_expt_dp_group=metadata.get('intra_expt_dp_group'), ) diff --git a/megatron/core/transformer/mlp.py b/megatron/core/transformer/mlp.py index ae0f8171fd4..cb3b1b0be82 100644 --- a/megatron/core/transformer/mlp.py +++ b/megatron/core/transformer/mlp.py @@ -178,6 +178,7 @@ def __init__( ffn_hidden_size: Optional[int] = None, tp_group: Optional[torch.distributed.ProcessGroup] = None, name: str | None = None, + pg_collection: Optional[ProcessGroupCollection] = None, ): """ Args: @@ -226,6 +227,7 @@ def __init__( is_expert=is_expert, tp_comm_buffer_name="fc1", tp_group=tp_group, + pg_collection=pg_collection, stride=fc1_stride, name=(name + ".linear_fc1") if name is not None else None, ) @@ -248,6 +250,7 @@ def __init__( is_expert=is_expert, tp_comm_buffer_name="fc2", tp_group=tp_group, + pg_collection=pg_collection, name=(name + ".linear_fc2") if name is not None else None, ) @@ -410,6 +413,7 @@ def as_mlp_submodule( config=config, submodules=submodules, tp_group=pg_collection.tp, + pg_collection=pg_collection, is_expert=is_expert, input_size=input_size, ffn_hidden_size=ffn_hidden_size, diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py index 558b1b07a15..c786b17137f 100644 --- a/megatron/core/transformer/module.py +++ b/megatron/core/transformer/module.py @@ -92,6 +92,8 @@ def sharded_state_dict( sharded_offsets=sharded_offsets, tp_group=tp_group, dp_cp_group=metadata['dp_cp_group'], + intra_dp_cp_group=metadata.get('intra_dp_cp_group'), + intra_expt_dp_group=metadata.get('intra_expt_dp_group'), ) # Recurse into submodules for name, module in self.named_children(): diff --git a/megatron/core/transformer/moe/shared_experts.py b/megatron/core/transformer/moe/shared_experts.py index 027d0a780ff..038a162f899 100644 --- a/megatron/core/transformer/moe/shared_experts.py +++ b/megatron/core/transformer/moe/shared_experts.py @@ -125,8 +125,13 @@ def __init__( "please set '--disable-bias-linear' instead." config.ffn_hidden_size = config.moe_shared_expert_intermediate_size - # TODO(Hepteract): pass pg_collection to MLP after refactoring MLP - super().__init__(config=config, submodules=submodules, tp_group=pg_collection.tp, name=name) + super().__init__( + config=config, + submodules=submodules, + tp_group=pg_collection.tp, + name=name, + pg_collection=pg_collection, + ) self.use_shared_expert_gate = gate if self.use_shared_expert_gate: diff --git a/megatron/core/transformer/utils.py b/megatron/core/transformer/utils.py index aee4e961b9e..84f248dee24 100644 --- a/megatron/core/transformer/utils.py +++ b/megatron/core/transformer/utils.py @@ -100,6 +100,8 @@ def make_sharded_tensors_for_checkpoint( extra_state_suffix: str = '_extra_state', tp_group: Optional[torch.distributed.ProcessGroup] = None, dp_cp_group: Optional[torch.distributed.ProcessGroup] = None, + intra_dp_cp_group: Optional[torch.distributed.ProcessGroup] = None, + intra_expt_dp_group: Optional[torch.distributed.ProcessGroup] = None, ): """Wraps tensors from transformer layers with ShardedTensor or ShardedObject. @@ -152,6 +154,8 @@ def make_sharded_tensors_for_checkpoint( extra_state_suffix=extra_state_suffix, tp_group=tp_group, dp_cp_group=dp_cp_group, + intra_dp_cp_group=intra_dp_cp_group, + intra_expt_dp_group=intra_expt_dp_group, ) sharded_state_dict = {} @@ -292,6 +296,8 @@ def sharded_state_dict_default( sharded_offsets, tp_group=tp_group, dp_cp_group=metadata['dp_cp_group'], + intra_dp_cp_group=metadata.get('intra_dp_cp_group'), + intra_expt_dp_group=metadata.get('intra_expt_dp_group'), ) return module_sharded_sd diff --git a/megatron/core/utils.py b/megatron/core/utils.py index 72373e9ac3b..1df7cf358c1 100644 --- a/megatron/core/utils.py +++ b/megatron/core/utils.py @@ -992,6 +992,8 @@ def make_tp_sharded_tensor_for_checkpoint( # Pop group parameters from kwargs tp_group = kwargs.pop('tp_group', None) dp_cp_group = kwargs.pop('dp_cp_group', None) + intra_dp_cp_group = kwargs.pop('intra_dp_cp_group', None) + intra_expt_dp_group = kwargs.pop('intra_expt_dp_group', None) prepend_axis_num = len(prepend_offsets) @@ -1051,10 +1053,17 @@ def make_tp_sharded_tensor_for_checkpoint( else: # GTP shards axis 0, TP shards a different axis → add a separate axis-0 offset new_offsets.append((prepend_axis_num, gtp_rank, gtp_remat_size)) - # Elect the writer over the gtp_remat-EXCLUDED DP group (its true replicas). - dp_replica_id = parallel_state.get_data_parallel_rank( - with_context_parallel=True, with_gtp_remat=False - ) + # Elect the writer over the GTP-remat-excluded replica group. Explicit-grid + # callers cannot use MPU globals, and expert weights replicate over expert DP. + replica_group = intra_dp_cp_group + if not getattr(tensor, 'allreduce', True) and intra_expt_dp_group is not None: + replica_group = intra_expt_dp_group + if replica_group is not None: + dp_replica_id = get_pg_rank(replica_group) + else: + dp_replica_id = parallel_state.get_data_parallel_rank( + with_context_parallel=True, with_gtp_remat=False + ) # Saved global is the padded shape when GTP padded out_features for alignment. if getattr(tensor, "pad_length", 0): kwargs.setdefault("allow_shape_mismatch", True) diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index 26650a7d559..3b5a842a799 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -2047,6 +2047,7 @@ def _maybe_setup_gpt_to_hybrid_load(args, ckpt_args, model): """ from megatron.core.dist_checkpointing.gpt_checkpoint_interop import gpt_compatible_layer_maps from megatron.core.models.hybrid.hybrid_model import HybridModel + from megatron.core.models.mimo.model import MimoModel def _contains_hybrid_model(module): # Megatron-FSDP and Float16Module both retain the wrapped module under @@ -2055,6 +2056,15 @@ def _contains_hybrid_model(module): while module is not None: if isinstance(module, HybridModel): return True + if isinstance(module, MimoModel): + if module.language_model is not None: + module = module.language_model + continue + # Heterogeneous encoder-only ranks do not instantiate the language + # model, but its spec still describes the model being checkpointed. + return bool( + module.mimo_config.language_model_spec.params.get('hybrid_layer_pattern') + ) module = getattr(module, 'module', None) return False diff --git a/megatron/training/initialize.py b/megatron/training/initialize.py index 87d6aa65b03..ed8cf3cdf6c 100644 --- a/megatron/training/initialize.py +++ b/megatron/training/initialize.py @@ -26,7 +26,13 @@ from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( enable_batch_invariant_mode, ) -from megatron.core.utils import get_pg_rank, get_te_version, is_te_min_version, is_torch_min_version +from megatron.core.utils import ( + get_pg_rank, + get_pg_size, + get_te_version, + is_te_min_version, + is_torch_min_version, +) from megatron.training import ( get_adlr_autoresume, get_args, @@ -426,10 +432,12 @@ def _set_random_seed( tp_group: Optional[torch.distributed.ProcessGroup] = None, ep_group: Optional[torch.distributed.ProcessGroup] = None, etp_group: Optional[torch.distributed.ProcessGroup] = None, + gtp_remat_group: Optional[torch.distributed.ProcessGroup] = None, + egtp_remat_group: Optional[torch.distributed.ProcessGroup] = None, ): """Set random seed for reproducability. - The optional pp/dp/tp/ep/etp groups let a caller without an initialized mpu + The optional parallel groups let a caller without an initialized mpu (e.g. a disjoint-grid run) supply the parallel ranks explicitly; each falls back to the mpu group when None. """ @@ -448,6 +456,16 @@ def _set_random_seed( tp_rank = get_pg_rank(tp_group) if tp_group is not None else None ep_rank = get_pg_rank(ep_group) if ep_group is not None else None etp_rank = get_pg_rank(etp_group) if etp_group is not None else None + gtp_remat_rank = get_pg_rank(gtp_remat_group) if gtp_remat_group is not None else None + egtp_remat_rank = ( + get_pg_rank(egtp_remat_group) if egtp_remat_group is not None else None + ) + gtp_remat_world_size = ( + get_pg_size(gtp_remat_group) if gtp_remat_group is not None else None + ) + egtp_remat_world_size = ( + get_pg_size(egtp_remat_group) if egtp_remat_group is not None else None + ) tensor_parallel.model_parallel_cuda_manual_seed( seed, te_rng_tracker, @@ -456,6 +474,10 @@ def _set_random_seed( tp_rank=tp_rank, ep_rank=ep_rank, etp_rank=etp_rank, + gtp_remat_rank=gtp_remat_rank, + egtp_remat_rank=egtp_remat_rank, + gtp_remat_world_size=gtp_remat_world_size, + egtp_remat_world_size=egtp_remat_world_size, ) else: raise ValueError("Seed ({}) should be a positive integer.".format(seed_)) diff --git a/megatron/training/training.py b/megatron/training/training.py index 8aab38e6071..abf6d4bd829 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -116,7 +116,6 @@ get_rerun_state_machine, ) from megatron.core.resharding.refit import swap_model_weights -from megatron.core.tensor_parallel.gtp_api import HAVE_GTP from megatron.core.transformer.cuda_graphs import TECudaGraphHelper from megatron.core.transformer.experimental_attention_variant.dsa import DSAIndexerLossLoggingHelper from megatron.core.transformer.module import Float16Module diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_custom_pgs.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_custom_pgs.py new file mode 100644 index 00000000000..a649c1ac36d --- /dev/null +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_custom_pgs.py @@ -0,0 +1,291 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""GTP_remat must follow the caller's ``pg_collection``, not the MPU globals. + +Two TransformerBlocks, same degrees (TP=1, CP=1, GTP_remat=2 over world=4), same weights, +same input: one built from ``parallel_state`` groups, one from a custom collection whose +``gtp_remat`` group is the PERMUTED pairing ([0,1],[2,3] vs [0,2],[1,3]). The forward +all-gathers every peer's shard, so both must produce identical output and gradients. + +The MPU globals stay initialized as the first topology throughout: a module that reads the +global group instead of the collection it was handed then gathers the wrong peer's shard -- +a valid-but-wrong group, which is the silent failure this test catches. +""" + +import pytest +import torch +import torch.distributed as dist + +from megatron.core.tensor_parallel.gtp_api import HAVE_GTP + +if not HAVE_GTP: + pytest.skip("GTP requires TransformerEngine >= 2.19", allow_module_level=True) + +from megatron.core.tensor_parallel.generalized_tensor_parallelism import GTPShardedParam +from tests.unit_tests.generalized_tensor_parallel.gtp_test_utils import ( + _requires_multi_gpu, + _run_distributed, + _torchrun_dist_init, + reset_fp8_state, + reset_gtp_globals, +) + +WORLD = 4 +GTP_SIZE = 2 +HIDDEN = 256 +NUM_HEADS = 8 +FFN_HIDDEN = 512 +NUM_LAYERS = 2 +SEQ = 16 +BATCH = 1 +dtype = torch.bfloat16 + +# The two ways to split a 4-rank world into gtp_remat pairs. Whichever one the MPU picks, +# the test uses the other. Every rank creates every group in this fixed order so the NCCL +# group tags agree across ranks -- a per-rank "create only the group I belong to" idiom +# assigns mismatched tags and hangs. +_PAIRINGS = {"adjacent": [[0, 1], [2, 3]], "strided": [[0, 2], [1, 3]]} + +# Forward is exact: all-gather is pure data movement, so both blocks feed bit-identical +# operands to identical GEMMs. Gradients additionally carry the attention backward's +# nondeterminism, hence the looser BF16-scale tolerance. +FWD_TOL = dict(atol=1e-5, rtol=1e-5) +GRAD_TOL = dict(atol=2e-2, rtol=2e-2) + + +def _make_config(): + from megatron.core.transformer.transformer_config import TransformerConfig + + return TransformerConfig( + num_attention_heads=NUM_HEADS, + num_layers=NUM_LAYERS, + hidden_size=HIDDEN, + ffn_hidden_size=FFN_HIDDEN, + add_bias_linear=False, + params_dtype=dtype, + hidden_dropout=0.0, + attention_dropout=0.0, + bias_dropout_fusion=False, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + ) + + +def _build_block(pg_collection): + """Build a GTP-sharded TransformerBlock wired to ``pg_collection``.""" + from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec + from megatron.core.transformer.transformer_block import TransformerBlock + + block = TransformerBlock( + _make_config(), get_gpt_layer_with_transformer_engine_spec(), pg_collection=pg_collection + ).cuda() + assert any( + isinstance(p, GTPShardedParam) for p in block.parameters() + ), "GTP is not active: the block has no GTPShardedParam" + return block + + +def _pick_permuted_gtp_group(rank, mpu_ranks): + """Create both candidate pairings on every rank; return this rank's group in the other one. + + Returns the group whose membership differs from ``mpu_ranks``, so any module that reads + the global group instead of the supplied one gathers a different peer's shard. + """ + my_groups = {} # sorted pair -> this rank's group in that pairing + for pairs in _PAIRINGS.values(): + for pair in pairs: + group = dist.new_group(ranks=pair) + if rank in pair: + my_groups[tuple(sorted(pair))] = group + + permuted = [g for pair, g in my_groups.items() if list(pair) != mpu_ranks] + assert len(permuted) == 1, ( + f"rank {rank}: want exactly one pairing differing from the MPU group {mpu_ranks}, " + f"got {list(my_groups)}" + ) + return permuted[0] + + +def _canonical_full_weights(block, gtp_group): + """Gather every parameter to full (unsharded) form, then broadcast rank 0's copy world-wide. + + Returns a name -> tensor dict that is bit-identical on every rank, so both blocks can be + loaded with the same global model no matter how the shards are distributed. + """ + full_weights = {} + for name, param in block.named_parameters(): + if isinstance(param, GTPShardedParam): + shards = [torch.empty_like(param.data) for _ in range(gtp_group.size())] + dist.all_gather(shards, param.data.contiguous(), group=gtp_group) + full = torch.cat(shards, dim=0) + else: + full = param.data.clone() + dist.broadcast(full, src=0) + full_weights[name] = full + return full_weights + + +def _load_full_weights(block, full_weights, gtp_rank): + """Load the canonical weights, slicing GTP params by ``gtp_rank`` and priming main_grad.""" + for name, param in block.named_parameters(): + full = full_weights[name] + if isinstance(param, GTPShardedParam): + shard = param.shape[0] + param.data.copy_(full[gtp_rank * shard : (gtp_rank + 1) * shard]) + # GTP writes the reduce-scattered wgrad here; it must exist before backward. + param.main_grad = torch.zeros(param.shape, dtype=dtype, device='cuda') + else: + param.data.copy_(full) + + +def _full_grads(block, gtp_group): + """Full (unsharded) gradients keyed by parameter name, for cross-topology comparison.""" + grads = {} + for name, param in block.named_parameters(): + if isinstance(param, GTPShardedParam): + shards = [torch.empty_like(param.main_grad) for _ in range(gtp_group.size())] + dist.all_gather(shards, param.main_grad.contiguous(), group=gtp_group) + grads[name] = torch.cat(shards, dim=0).float().cpu() + elif param.grad is not None: + grads[name] = param.grad.detach().float().cpu() + return grads + + +def _fwd_bwd(block, x): + """Run one forward/backward; return (output, input gradient) on cpu in fp32.""" + out = block(hidden_states=x, attention_mask=None) + out.sum().backward() + return out.detach().float().cpu(), x.grad.detach().float().cpu() + + +def _worker_custom_pgs_match_mpu(rank, world_size, port): + from megatron.core import parallel_state as ps + from megatron.core.process_groups_config import ProcessGroupCollection + from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + + # ---------------- Topology 1: groups from parallel_state ---------------- + ps.destroy_model_parallel() + ps.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=GTP_SIZE + ) + model_parallel_cuda_manual_seed(42) + + mpu_pgs = ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['tp', 'cp', 'pp', 'gtp_remat', 'expt_gtp_remat'] + ) + mpu_gtp_group = mpu_pgs.gtp_remat + assert ( + mpu_gtp_group.size() == GTP_SIZE + ), f"GTP_remat inactive: group size {mpu_gtp_group.size()}, want {GTP_SIZE}" + + block_mpu = _build_block(mpu_pgs) + + # One canonical global model, shared by both topologies. + full_weights = _canonical_full_weights(block_mpu, mpu_gtp_group) + _load_full_weights(block_mpu, full_weights, mpu_gtp_group.rank()) + + torch.manual_seed(1234) + x = torch.randn(SEQ, BATCH, HIDDEN, dtype=dtype, device='cuda') + dist.broadcast(x, src=0) # identical input on every rank + + out_mpu, grad_in_mpu = _fwd_bwd(block_mpu, x.clone().requires_grad_(True)) + grads_mpu = _full_grads(block_mpu, mpu_gtp_group) + + del block_mpu + GTPShardedParam._chain_state = {} + + # ---------------- Topology 2: custom collection, permuted gtp ranks ---------------- + mpu_ranks = sorted(dist.get_process_group_ranks(mpu_gtp_group)) + custom_gtp_group = _pick_permuted_gtp_group(rank, mpu_ranks) + + # Only gtp_remat differs; tp/cp/pp are size-1 groups, identical in both topologies. + custom_pgs = ProcessGroupCollection( + tp=mpu_pgs.tp, + cp=mpu_pgs.cp, + pp=mpu_pgs.pp, + gtp_remat=custom_gtp_group, + expt_gtp_remat=mpu_pgs.expt_gtp_remat, + ) + # Seed from the custom topology's gtp rank rather than the global one. The weights are + # overwritten below, so this only has to be self-consistent -- it also exercises the + # explicit-rank arguments of model_parallel_cuda_manual_seed. + model_parallel_cuda_manual_seed( + 42, gtp_remat_rank=custom_gtp_group.rank(), egtp_remat_rank=0, force_reset_rng=True + ) + + block_custom = _build_block(custom_pgs) + _load_full_weights(block_custom, full_weights, custom_gtp_group.rank()) + + out_custom, grad_in_custom = _fwd_bwd(block_custom, x.clone().requires_grad_(True)) + grads_custom = _full_grads(block_custom, custom_gtp_group) + + ps.destroy_model_parallel() + ps.initialize_model_parallel() + + # ---------------- The two topologies must agree ---------------- + torch.testing.assert_close( + out_custom, + out_mpu, + **FWD_TOL, + msg="forward output differs between MPU and custom gtp_remat groups", + ) + torch.testing.assert_close( + grad_in_custom, + grad_in_mpu, + **GRAD_TOL, + msg="input gradient differs between MPU and custom gtp_remat groups", + ) + assert set(grads_custom) == set(grads_mpu), "parameter sets differ between the two blocks" + for name in sorted(grads_mpu): + torch.testing.assert_close( + grads_custom[name], + grads_mpu[name], + **GRAD_TOL, + msg=f"weight gradient for {name} differs between MPU and custom gtp_remat groups", + ) + + +def _worker_partial_pgs_fall_back_to_mpu(rank, world_size, port): + """A collection that omits gtp_remat must fall back to the MPU group, not disable GTP. + + ``__getattr__`` returns None for unset fields, so ``hasattr`` lies: a resolver trusting it + reads None and silently builds an unsharded block. + """ + from megatron.core import parallel_state as ps + from megatron.core.process_groups_config import ProcessGroupCollection + from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + + ps.destroy_model_parallel() + ps.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=GTP_SIZE + ) + model_parallel_cuda_manual_seed(42) + + partial_pgs = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp', 'pp']) + assert 'gtp_remat' not in vars(partial_pgs), "this collection must omit gtp_remat" + + block = _build_block(partial_pgs) + + gtp_group = ps.get_gtp_weight_remat_group() + sharded = [(n, p) for n, p in block.named_parameters() if isinstance(p, GTPShardedParam)] + assert sharded, "no parameter was sharded: the resolver did not fall back to the MPU group" + for name, param in sharded: + assert param.gtp_remat_size == gtp_group.size(), ( + f"{name} was sharded over a size-{param.gtp_remat_size} axis, " + f"want {gtp_group.size()} (the MPU gtp_remat group)" + ) + + ps.destroy_model_parallel() + ps.initialize_model_parallel() + + +class TestGTPCustomProcessGroups: + def test_custom_gtp_pg_collection_matches_mpu(self): + """A permuted-but-equivalent gtp_remat group must give identical fwd/bwd results.""" + _requires_multi_gpu(WORLD) + _run_distributed(_worker_custom_pgs_match_mpu, WORLD) + + def test_pg_collection_without_gtp_remat_falls_back_to_mpu(self): + """Omitting gtp_remat must fall back to the MPU group, not silently disable sharding.""" + _requires_multi_gpu(WORLD) + _run_distributed(_worker_partial_pgs_fall_back_to_mpu, WORLD) diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_grad_correctness.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_grad_correctness.py index b969e5e5bf8..1440ba137fa 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_gtp_grad_correctness.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_grad_correctness.py @@ -109,12 +109,16 @@ def _run_one_backward(ddp_model, rank, calculate_per_token_loss=False): # overlap_grad_reduce=False. Do NOT also call start_grad_sync() — that double- # reduces, which is idempotent at full-DP size but halves at replicate size. ddp_model.finish_grad_sync() + from megatron.core import parallel_state as ps from megatron.core.distributed.finalize_model_grads import ( _allreduce_replicated_grads_over_gtp_remat_group, ) _allreduce_replicated_grads_over_gtp_remat_group( - [ddp_model], calculate_per_token_loss=calculate_per_token_loss + [ddp_model], + ps.get_gtp_weight_remat_group(check_initialized=False), + ps.get_expert_gtp_weight_remat_group(check_initialized=False), + calculate_per_token_loss=calculate_per_token_loss, ) return float(loss.item()) @@ -274,11 +278,16 @@ def _run_step_distopt(ddp_model, optim, rank): loss.backward() # Production order (finalize_model_grads): reduce across DP first, THEN the gtp_remat finalize. ddp_model.finish_grad_sync() + from megatron.core import parallel_state as ps from megatron.core.distributed.finalize_model_grads import ( _allreduce_replicated_grads_over_gtp_remat_group, ) - _allreduce_replicated_grads_over_gtp_remat_group([ddp_model]) + _allreduce_replicated_grads_over_gtp_remat_group( + [ddp_model], + ps.get_gtp_weight_remat_group(check_initialized=False), + ps.get_expert_gtp_weight_remat_group(check_initialized=False), + ) _, grad_norm, _ = optim.step() return float(grad_norm) diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_muon_dcp.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_muon_dcp.py index b26d8a974ce..1dff77998c7 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_gtp_muon_dcp.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_muon_dcp.py @@ -125,7 +125,6 @@ def test_gtp_muon_moe_save_load(self, tmp_path_dist_ckpt): if int(os.environ.get('WORLD_SIZE', '1')) != 4: pytest.skip("Requires world_size 4 (gtp2 x dp2)") - os.environ['MEGATRON_GTP_FORCE_ENABLE'] = '1' from megatron.core import parallel_state as ps from megatron.core.tensor_parallel import model_parallel_cuda_manual_seed from megatron.core.tensor_parallel.generalized_tensor_parallelism import ( @@ -237,7 +236,6 @@ def test_gtp_muon_moe_native_fp8_save_load(self, tmp_path_dist_ckpt): pytest.skip("Requires world_size 4 (gtp2 x dp2)") _requires_mxfp8() - os.environ['MEGATRON_GTP_FORCE_ENABLE'] = '1' from megatron.core import parallel_state as ps from megatron.core.fp8_utils import is_float8tensor from megatron.core.tensor_parallel import model_parallel_cuda_manual_seed diff --git a/tests/unit_tests/models/mimo/test_mimo_hetero_e2e_train_checkpoint.py b/tests/unit_tests/models/mimo/test_mimo_hetero_e2e_train_checkpoint.py index 91daecb1706..783a73f3fda 100644 --- a/tests/unit_tests/models/mimo/test_mimo_hetero_e2e_train_checkpoint.py +++ b/tests/unit_tests/models/mimo/test_mimo_hetero_e2e_train_checkpoint.py @@ -1,6 +1,6 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -"""End-to-end: the hetero MIMO 20L mock trains and round-trips a checkpoint. +"""End-to-end: the hetero MIMO 20L GTP mock trains and round-trips a checkpoint. This drives the training launcher, which spawns its own 8-rank ``torch.distributed.run``, so it must run as a single plain pytest process (not under the multi-rank unit runner). @@ -30,6 +30,10 @@ def _run_launcher(base, train_iters, extra_args, name): **os.environ, "TRAIN_ITERS": str(train_iters), "TORCHRUN_LOG_DIR": str(base / f"torchrun-{name}"), + "LLM_DP": "1", + "LLM_EP": "2", + "TENSOR_PARALLEL_NUM_WEIGHT_SHARDS": "4", + "EXPERT_TENSOR_PARALLEL_NUM_WEIGHT_SHARDS": "2", } # conftest's autouse set_env fixture disables TE flash/fused attention; the 20L model # at seq 8192 needs them (unfused attention OOMs), so let the launcher use TE defaults. @@ -63,7 +67,7 @@ def _tail(result): @pytest.mark.skipif( _UNDER_TORCHRUN, reason="launcher spawns its own torchrun; run as a plain process" ) -def test_hetero_mimo_20l_trains_and_checkpoint_round_trips(): +def test_hetero_mimo_20l_gtp_trains_and_checkpoint_round_trips(): # The 128-expert MoE checkpoint is large; save under the repo workspace (a roomy # shared filesystem on the cluster) rather than pytest's node-local /tmp tmp_path. scratch = Path(tempfile.mkdtemp(prefix="mimo_e2e_", dir=_REPO_ROOT)) diff --git a/tests/unit_tests/models/mimo/test_mimo_hetero_grid_args.py b/tests/unit_tests/models/mimo/test_mimo_hetero_grid_args.py index fb923c2b37d..cee45406b39 100644 --- a/tests/unit_tests/models/mimo/test_mimo_hetero_grid_args.py +++ b/tests/unit_tests/models/mimo/test_mimo_hetero_grid_args.py @@ -62,6 +62,33 @@ def test_canonical_layout_validates_and_maps_specs(): assert language_grid_spec.expt_tp == 1 +def test_gtp_layout_validates_and_maps_weight_shard_axes(): + args = _layout_8gpu_20l( + llm_dp=1, + llm_ep=2, + tensor_parallel_num_weight_shards=4, + expert_tensor_parallel_num_weight_shards=2, + ) + + assert validate_hetero_grid_args(args, WORLD_SIZE_8) == (4, 4) + assert args.gtp_weight_remat_size == 2 + assert args.expert_gtp_weight_remat_size == 2 + + _, language_grid_spec = build_module_grid_specs( + args, WORLD_SIZE_8, encoder_module_name="radio_encoder" + ) + assert language_grid_spec.gtp_remat == 2 + assert language_grid_spec.dp == 1 + assert language_grid_spec.expt_gtp_remat == 2 + assert language_grid_spec.expt_dp == 1 + + +def test_weight_shards_must_divide_language_tp(): + args = _layout_8gpu_20l(tensor_parallel_num_weight_shards=3) + with pytest.raises(ValueError, match="must be divisible"): + validate_hetero_grid_args(args, WORLD_SIZE_8) + + def test_overlapping_spans_raise(): # llm-offset 2 makes llm ranks {2,3,4,5} overlap encoder ranks {0,1,2,3}. args = _layout_8gpu_20l(llm_offset=2) diff --git a/tests/unit_tests/pipeline_parallel/test_bridge_communicator.py b/tests/unit_tests/pipeline_parallel/test_bridge_communicator.py index e4801ad8939..4eaa11b8c24 100644 --- a/tests/unit_tests/pipeline_parallel/test_bridge_communicator.py +++ b/tests/unit_tests/pipeline_parallel/test_bridge_communicator.py @@ -112,7 +112,7 @@ def _shard_and_copy_( _active_grids: list = [] -def create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1): +def create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1, gtp_remat=1): """Create a HyperCommGrid with tensor parallelism=2, context parallelism=2, and data parallelism=2.""" # Set up environment for world size 8 if not already set if not dist.is_initialized(): @@ -123,12 +123,13 @@ def create_hypercomm_grid(offset=0, tp=1, cp=1, pp=1, dp=1): os.environ["WORLD_SIZE"] = "8" grid = HyperCommGrid( - shape=[tp, cp, pp, dp], - dim_names=["tp", "cp", "pp", "dp"], + shape=[tp, gtp_remat, cp, pp, dp], + dim_names=["tp", "gtp_remat", "cp", "pp", "dp"], rank_offset=offset, backend="nccl", ) _ = grid.create_pg(["tp"]) + _ = grid.create_pg(["gtp_remat"]) _ = grid.create_pg(["cp"]) _ = grid.create_pg(["pp"]) _ = grid.create_pg(["dp"]) @@ -326,6 +327,17 @@ def test_bridge_pg_membership(self, grid1_tp, grid1_dp, grid2_tp, grid2_dp): ] assert all(rank not in expected for rank in member_ranks) + def test_gtp_is_an_independent_bridge_data_lane(self): + src_grid = create_hypercomm_grid(offset=0, tp=2, dp=2) + dest_grid = create_hypercomm_grid(offset=4, tp=2, dp=1, gtp_remat=2) + bridge = BridgeCommunicator(src_grid, dest_grid) + + assert len(bridge.src_tp_leaders) == 2 + assert len(bridge.dest_tp_leaders) == 2 + assert sorted(set(bridge.src_tp_leaders) | set(bridge.dest_tp_leaders)) == list( + dist.get_process_group_ranks(bridge.bridge_pg) + ) + def test_send_forward_recv_forward(self): """Test send_forward and recv_forward operations.""" diff --git a/tests/unit_tests/test_checkpointing.py b/tests/unit_tests/test_checkpointing.py index cf1b6a76539..6fa9599f3d2 100644 --- a/tests/unit_tests/test_checkpointing.py +++ b/tests/unit_tests/test_checkpointing.py @@ -11,6 +11,7 @@ from megatron.core.distributed import DistributedDataParallelConfig from megatron.core.distributed.fsdp.mcore_fsdp_adapter import FullyShardedDataParallel +from megatron.core.models.mimo.model import MimoModel from megatron.core.num_microbatches_calculator import ( init_num_microbatches_calculator, unset_num_microbatches_calculator, @@ -23,6 +24,7 @@ CheckpointType, _build_sharded_state_dict_metadata, _load_base_checkpoint, + _maybe_setup_gpt_to_hybrid_load, get_checkpoint_tracker_filename, load_checkpoint, maybe_save_dataloader_state, @@ -75,6 +77,22 @@ def sharded_state_dict(self, *args, metadata: Optional[dict] = None, **kwargs): return self.state_dict() +def test_hybrid_checkpoint_matches_encoder_only_mimo_rank(): + """An encoder-only MIMO rank identifies its uninstantiated hybrid language model.""" + model = MimoModel.__new__(MimoModel) + torch.nn.Module.__init__(model) + model.language_model = None + model.mimo_config = SimpleNamespace( + language_model_spec=SimpleNamespace(params={'hybrid_layer_pattern': 'M*'}) + ) + + assert _maybe_setup_gpt_to_hybrid_load( + SimpleNamespace(hybrid_layer_pattern='M*'), + SimpleNamespace(hybrid_layer_pattern='M*'), + [model], + ) == (None, False) + + def test_maybe_save_dataloader_state_uses_explicit_process_groups(tmp_path): """Dataloader checkpoints use the supplied module groups and canonical model-parallel path.""" groups = { diff --git a/tests/unit_tests/test_mimo_hetero_topology.py b/tests/unit_tests/test_mimo_hetero_topology.py index 7d68e107bf7..95ba6c5c7d4 100644 --- a/tests/unit_tests/test_mimo_hetero_topology.py +++ b/tests/unit_tests/test_mimo_hetero_topology.py @@ -25,6 +25,21 @@ def _specs(): ] +def _gtp_specs(): + return [ + ModuleGridSpec(name=ENCODER, num_ranks=4, tp=2, rank_offset=0), + ModuleGridSpec( + name=MIMO_LANGUAGE_MODULE_KEY, + num_ranks=4, + tp=2, + gtp_remat=2, + ep=2, + expt_gtp_remat=2, + rank_offset=4, + ), + ] + + class TestModuleGridSpecResolution: def test_derived_dims_resolve_to_concrete_ints(self): # num_ranks=4,tp=2 with default expt_tp=1: dp=2, expt_dp=4. @@ -87,6 +102,33 @@ def test_pgc_group_sizes(self): finally: topo.destroy() + def test_gtp_pgc_group_sizes(self): + topo = create_topology(_gtp_specs()) + try: + rank = dist.get_rank() + pgc = ( + topo.module_pgs[ENCODER] + if rank < 4 + else topo.module_pgs[MIMO_LANGUAGE_MODULE_KEY] + ) + if rank < 4: + assert pgc.gtp_remat.size() == 1 + assert pgc.expt_gtp_remat.size() == 1 + else: + assert pgc.tp.size() == 2 + assert pgc.gtp_remat.size() == 2 + assert pgc.dp.size() == 1 + assert pgc.dp_cp_gtp_remat.size() == 2 + assert pgc.mp.size() == 4 + assert pgc.expt_tp.size() == 1 + assert pgc.ep.size() == 2 + assert pgc.expt_gtp_remat.size() == 2 + assert pgc.expt_dp.size() == 1 + assert pgc.expt_dp_gtp_remat.size() == 2 + assert pgc.tp_ep_pp_with_egtp_remat.size() == 4 + finally: + topo.destroy() + def test_embedding_groups(self): # Language grid is tp=2,pp=2 at ranks 4-7: each PP group is [first,last] (size 2), # so first/last-stage ranks get a 2-rank .embd and the first stage gets .pos_embd.