From 8cd3004d48de4702bc2ab8a47bc69d45c9985ad2 Mon Sep 17 00:00:00 2001 From: Shiqing Fan Date: Sun, 30 Aug 2026 22:00:55 -0700 Subject: [PATCH] (feat) gtp+cpu-init support - Support CPU initialization with GTP weight rematerialization - Refresh the quantized-weight init snapshot after CPU initialization - Replicate the GTP linear bias instead of sharding it - Add corresponding UTs for above 3 cases Signed-off-by: Shiqing Fan --- .../generalized_tensor_parallelism.py | 51 ++++ megatron/core/tensor_parallel/gtp_api.py | 2 + megatron/core/tensor_parallel/layers.py | 13 +- .../test_gtp_basics.py | 81 ++++++ .../test_gtp_cpu_init.py | 253 ++++++++++++++++++ 5 files changed, 399 insertions(+), 1 deletion(-) create mode 100644 tests/unit_tests/generalized_tensor_parallel/test_gtp_cpu_init.py diff --git a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py b/megatron/core/tensor_parallel/generalized_tensor_parallelism.py index 6e1d4ba275b..00e72f6caa1 100644 --- a/megatron/core/tensor_parallel/generalized_tensor_parallelism.py +++ b/megatron/core/tensor_parallel/generalized_tensor_parallelism.py @@ -551,6 +551,26 @@ def gtp_remat_shard_dim0(dim0, gtp_remat_group): return padded // gtp_remat_size, pad_length +def gtp_remat_slice_rows(unsharded, gtp_remat_group): + """Take this rank's GTP_remat rows out of a tensor that is not yet GTP-sharded. + + Use this to fill a weight that was **already allocated pre-sharded**; padding matches + :func:`gtp_remat_shard_dim0`, so the rows line up with that shard. Unlike + :func:`_gtp_slice_one_param`, this returns a plain tensor and replaces nothing. + + Returns: + This rank's rows, sized like one GTP shard (padding included). + """ + shard_dim0, pad_length = gtp_remat_shard_dim0(unsharded.shape[0], gtp_remat_group) + if pad_length > 0: + # Append the alignment rows to the tail of dim 0. Spelled as a concat rather than F.pad, + # whose spec is ordered last-dim-first and so reads backwards for a dim-0 pad. + alignment_rows = unsharded.new_zeros(pad_length, *unsharded.shape[1:]) + unsharded = torch.cat([unsharded, alignment_rows], dim=0) + gtp_rank = gtp_remat_group.rank() + return unsharded[gtp_rank * shard_dim0 : (gtp_rank + 1) * shard_dim0] + + def _gtp_slice_one_param(param, gtp_remat_group, *, name=""): """Pad + slice a full-size BF16 weight to this rank's GTP shard. @@ -714,6 +734,37 @@ def attach_gtp_to_presharded_module( new_weights.append(gtp_param) if is_grouped and new_weights: new_weights[0].weight_list = new_weights + _gtp_restore_replicated_bias(module, gtp_remat_group, pad_length, is_grouped=is_grouped) + + +def _gtp_restore_replicated_bias(module, gtp_remat_group, pad_length, is_grouped=False): + """Restore the module's linear bias(es) to full size, replicated across the GTP group. + + ``_gtp_pre_init`` passes TE a pre-sharded ``out_features`` so that TE builds the weight already + sharded; TE sizes the bias from that same value, so the bias comes out sharded as well. GTP + shards weights only, so this undoes that side effect. + + Re-allocating zeros reproduces the values exactly because TE zero-initializes every linear bias + (``init_method_constant(0.0)``); a future non-zero bias init would need a gather instead. + """ + # A packed bias has no per-gemm entries to re-allocate, so it would silently stay sharded. + assert not getattr( + module, "single_grouped_bias", False + ), f"GTP grouped module {type(module).__name__} requires single_grouped_bias=False." + # GroupedLinear does not declare bias_names; its biases are bias0..bias{num_gemms-1}. + bias_names = getattr(module, "bias_names", None) + if not bias_names: + bias_names = [f"bias{idx}" for idx in range(module.num_gemms)] if is_grouped else ["bias"] + gtp_remat_size = gtp_remat_group.size() + for name in bias_names: + bias = getattr(module, name, None) + # TE registers an empty plain tensor (not a Parameter) for the bias-free case. + if not isinstance(bias, torch.nn.Parameter) or bias.numel() == 0: + continue + # Resize in place: the Parameter keeps its identity, so its module registration and the + # attributes DDP / the optimizer / checkpointing read off it all survive untouched. + with torch.no_grad(): + bias.data = bias.new_zeros(bias.shape[0] * gtp_remat_size - pad_length) # Cache of dynamic ``GTP_`` subclasses, keyed by the FP8 base class. diff --git a/megatron/core/tensor_parallel/gtp_api.py b/megatron/core/tensor_parallel/gtp_api.py index 0e0a7e67725..b3be8847880 100644 --- a/megatron/core/tensor_parallel/gtp_api.py +++ b/megatron/core/tensor_parallel/gtp_api.py @@ -23,6 +23,7 @@ get_rs_stream, gtp_native_fp8_load_context, gtp_remat_shard_dim0, + gtp_remat_slice_rows, gtp_replica_rank, initialize_graph_wgrad_rings, is_gtp_param, @@ -67,6 +68,7 @@ def deregister_and_clear_gtp_symm_pools() -> None: "get_rs_stream", "gtp_native_fp8_load_context", "gtp_remat_shard_dim0", + "gtp_remat_slice_rows", "gtp_replica_rank", "is_gtp_param", "initialize_graph_wgrad_rings", diff --git a/megatron/core/tensor_parallel/layers.py b/megatron/core/tensor_parallel/layers.py index b77c4d202c0..7b3b8e1daa1 100644 --- a/megatron/core/tensor_parallel/layers.py +++ b/megatron/core/tensor_parallel/layers.py @@ -255,7 +255,8 @@ def _initialize_affine_weight_cpu( """Initialize affine weight for model parallel. Build the master weight on all processes and scatter - the relevant chunk.""" + the relevant chunk. A ``weight`` that is already GTP_remat-sharded is sliced down to this + rank's GTP rows as well, so the initialization matches a run with GTP off.""" if not skip_set_tensor_parallel_attributes: set_tensor_model_parallel_attributes( @@ -277,7 +278,17 @@ def _initialize_affine_weight_cpu( with torch.no_grad(): # all tensors must live on the same device cpu_weight = torch.cat(my_weight_list, dim=partition_dim).to_dense() + if getattr(weight, "gtp_remat_size", 1) > 1: + from megatron.core.tensor_parallel.gtp_api import gtp_remat_slice_rows + + cpu_weight = gtp_remat_slice_rows(cpu_weight, weight.group) weight.data.copy_(cpu_weight) + # Quantized (FP8/FP4) primary weights snapshot their values at construction, before CPU + # init writes the real ones. The distributed optimizer seeds its master params from that + # snapshot, so it must be refreshed or training starts from uninitialized memory. + high_precision_init_val = getattr(weight, "_high_precision_init_val", None) + if high_precision_init_val is not None: + high_precision_init_val.copy_(cpu_weight) if return_master_weight: return master_weight return None diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py index fb9c4a9a0dd..a499164177a 100644 --- a/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py @@ -1921,3 +1921,84 @@ def test_count_zeros_excludes_gtp_padding_end_to_end(self): .gtp_pad_zeros wiring.""" _requires_multi_gpu(4) _run_distributed(_worker_count_zeros_excludes_gtp_padding, 4) + + +# --------------------------------------------------------------------------- +# TestGTPReplicatedBias - the linear bias must stay replicated, not follow the weight +# --------------------------------------------------------------------------- + + +def _worker_bias_is_replicated(rank, world_size, port, gtp_remat_size): + from megatron.core import parallel_state as ps + from megatron.core.extensions.transformer_engine import TEColumnParallelLinear + from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + from megatron.core.transformer.transformer_config import TransformerConfig + + hidden, out_features, bias_value = 256, 512, 100.0 + ps.destroy_model_parallel() + ps.initialize_model_parallel(gtp_remat_size=gtp_remat_size) + try: + model_parallel_cuda_manual_seed(1234, force_reset_rng=True) + pg_collection = ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['tp', 'cp', 'pp', 'ep', 'expt_tp', 'gtp_remat', 'expt_gtp_remat'] + ) + layer = TEColumnParallelLinear( + hidden, + out_features, + config=TransformerConfig( + num_layers=1, + hidden_size=hidden, + num_attention_heads=8, + add_bias_linear=True, + params_dtype=torch.bfloat16, + bf16=True, + ), + init_method=lambda w: nn.init.normal_(w, std=0.02), + gather_output=False, + bias=True, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name=None, + pg_collection=pg_collection, + ).cuda() + + assert layer.bias.shape == (out_features,), ( + f"bias is {tuple(layer.bias.shape)} at GTP_remat={gtp_remat_size}; expected the " + f"replicated size ({out_features},)" + ) + + # The shape alone is not enough: a sharded bias fails silently, still returning the full + # out_features but biasing only the first out_features/gtp_remat_size columns. + with torch.no_grad(): + layer.bias.fill_(bias_value) + shard = layer.weight.data.contiguous() + shards = [torch.empty_like(shard) for _ in range(gtp_remat_size)] + dist.all_gather(shards, shard, group=pg_collection.gtp_remat) + full_weight = torch.cat(shards, dim=0)[:out_features].float() + + inputs = torch.randn(8, hidden, dtype=torch.bfloat16, device="cuda") + dist.broadcast(inputs, src=0) + # GTP writes the reduce-scattered wgrad here; it must exist before backward. + layer.weight.main_grad = torch.zeros( + layer.weight.shape, dtype=torch.bfloat16, device="cuda" + ) + + applied = (layer(inputs)[0].float() - inputs.float() @ full_weight.T).mean(dim=0) + missing = (applied - bias_value).abs() > 1.0 + assert not missing.any(), ( + f"{int(missing.sum())} of {out_features} outputs did not receive the bias at " + f"GTP_remat={gtp_remat_size}: the bias is still sharded" + ) + finally: + ps.destroy_model_parallel() + GTPShardedParam._chain_state = {} + + +class TestGTPReplicatedBias: + @pytest.mark.parametrize("gtp_remat_size", [2, 4]) + def test_bias_is_replicated(self, gtp_remat_size): + """``_gtp_pre_init`` gives TE a pre-sharded ``out_features`` and TE sizes the bias from it, + so without ``_gtp_restore_replicated_bias`` the bias is sharded too (256 of 512 outputs + biased at GTP_remat=2, with no error).""" + _requires_multi_gpu(4) + _run_distributed(_worker_bias_is_replicated, 4, gtp_remat_size) diff --git a/tests/unit_tests/generalized_tensor_parallel/test_gtp_cpu_init.py b/tests/unit_tests/generalized_tensor_parallel/test_gtp_cpu_init.py new file mode 100644 index 00000000000..dedbef6563f --- /dev/null +++ b/tests/unit_tests/generalized_tensor_parallel/test_gtp_cpu_init.py @@ -0,0 +1,253 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""CPU initialization must produce the same global model at every GTP_remat degree. + +CPU init exists so a run's initial weights are a function of the seed alone, not of the parallel +layout -- that is what makes convergence comparisons across configurations meaningful. The TE +modules allocate their weight *pre-sharded* under GTP_remat (see ``_gtp_pre_init``), so the CPU path +has to pad-and-slice the master weight the same way (see ``_initialize_affine_weight_cpu``); +otherwise it either shape-mismatches or hands every GTP peer the same rows. + +Everything below is exercised on ONE world_size=4, TP=1 stack holding every layer kind at once -- +GDP mixer, attention, MLP: + +* **GTP_remat 1 / 2 / 4**, each compared **bitwise** against the GTP_remat=1 reference. +* **Both TE parallel modes** -- column (``in_proj``, ``linear_qkv``, ``fc1``) and row + (``out_proj``, ``linear_proj``, ``fc2``), which are the two slicing geometries. +* **Alignment padding.** The GDP ``in_proj`` is 3104 wide, not a multiple of + ``32 * gtp_remat_size`` at any degree here, so the padded branch of ``gtp_remat_slice_rows`` + always runs -- asserted rather than assumed. +* **BF16 and MXFP8 primary weights** (``fp8_param``). +* **The quantized-init snapshot.** TE caches ``_high_precision_init_val`` at construction and the + distributed optimizer seeds its fp32 master params from it, but CPU init writes the real values + only afterwards -- a stale snapshot means training starts from uninitialized memory. Checked at + GTP_remat=1 as well, because that bug needs no GTP. +* **Replicated params stay replicated** -- layernorms, and the GDP conv1d / A_log / dt_bias, which + are GPU-initialized regardless and must come through the CPU path untouched. +""" + +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.gtp_api import dequantize_gtp_native_fp8, is_gtp_param +from tests.unit_tests.generalized_tensor_parallel.gtp_test_utils import ( + _requires_multi_gpu, + _requires_mxfp8, + _run_distributed, + _torchrun_dist_init, + reset_fp8_state, + reset_gtp_globals, +) + +WORLD = 4 +# GTP_remat degrees exercised; WORLD must be divisible by each. +GTP_SIZES = [2, 4] +HIDDEN = 256 +NUM_HEADS = 8 +FFN_HIDDEN = 512 +# One stack holding every layer kind at once: GDP mixer, attention, MLP. +LAYER_PATTERN = ["M", "*", "-"] +NUM_LAYERS = len(LAYER_PATTERN) +SEED = 1234 +dtype = torch.bfloat16 + + +def _make_config(fp8_param): + from megatron.core.transformer.transformer_config import TransformerConfig + + fp8_kwargs = dict(fp8="hybrid", fp8_recipe="mxfp8", fp8_param=True) if fp8_param else {} + 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, + bf16=True, + hidden_dropout=0.0, + attention_dropout=0.0, + bias_dropout_fusion=False, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + use_cpu_initialization=True, + is_hybrid_model=True, + mamba_num_groups=4, + mamba_num_heads=8, + mamba_head_dim=32, + **fp8_kwargs, + ) + + +def _build_block(config, pg_collection): + """Build a CPU-initialized GDP + attention + MLP stack and move it to the GPU.""" + from megatron.core.fp8_utils import get_fp8_context + from megatron.core.models.hybrid.hybrid_block import HybridStack + from megatron.core.models.hybrid.hybrid_layer_specs import gdp_stack_spec + + # Layers are normally constructed inside this context (transformer_block.py); it is what + # turns fp8_param into quantized primary weights. + with get_fp8_context(config, is_init=True): + block = HybridStack( + config, + gdp_stack_spec.submodules, + layer_type_list=LAYER_PATTERN, + pg_collection=pg_collection, + ) + return block.cuda() + + +def _local_values(param): + """This rank's parameter values as fp32 on cpu, dequantizing quantized primary weights.""" + if getattr(param, "_gtp_native_fp8", False): + # GTP reclasses native-FP8 shards into a dynamic subclass that TE's dequantize rejects. + data = dequantize_gtp_native_fp8(param) + elif hasattr(param, "dequantize"): + data = param.dequantize() + else: + data = param.data + return data.detach().float().cpu() + + +def _full_weights(block, gtp_group): + """name -> full (un-sharded, un-padded) fp32 values, with the padded tail dropped.""" + full = {} + for name, param in block.named_parameters(): + values = _local_values(param) + if is_gtp_param(param): + local = values.cuda() + shards = [torch.empty_like(local) for _ in range(gtp_group.size())] + dist.all_gather(shards, local, group=gtp_group) + values = torch.cat([s.cpu() for s in shards], dim=0) + pad_length = getattr(param, "pad_length", 0) + if pad_length: + values = values[: values.shape[0] - pad_length] + full[name] = values + return full + + +def _init_mpu(gtp_remat_size): + from megatron.core import parallel_state as ps + from megatron.core.process_groups_config import ProcessGroupCollection + from megatron.core.tensor_parallel.gtp_api import configure_gtp_remat_from_recipe + from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + + # Training calls this before model construction. Padding is enabled unconditionally (not just + # for the fp8 cases) so every build shards identically, and so the padded branch of + # gtp_remat_slice_rows is exercised: the GDP in_proj is 3104 wide, which is not a multiple of + # 32 * gtp_remat_size at any degree under test. + configure_gtp_remat_from_recipe(fp8=True, fp8_recipe="mxfp8") + ps.destroy_model_parallel() + ps.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1, gtp_remat_size=gtp_remat_size + ) + # CPU initialization draws from the CPU generator, which every rank seeds identically; the + # CUDA tracker is seeded too so a layer that falls back to GPU init is not silently skipped. + torch.manual_seed(SEED) + model_parallel_cuda_manual_seed(SEED, force_reset_rng=True) + return ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['tp', 'cp', 'pp', 'gtp_remat', 'expt_gtp_remat'] + ) + + +def _worker_cpu_init_is_gtp_invariant(rank, world_size, port, fp8_param, gtp_remat_size): + from megatron.core import parallel_state as ps + + config = _make_config(fp8_param) + + # ---------------- Reference: GTP_remat off, every rank holds the full weight ------------ + ref_pgs = _init_mpu(gtp_remat_size=1) + ref_block = _build_block(config, ref_pgs) + assert not any( + is_gtp_param(p) for p in ref_block.parameters() + ), "reference block must not be GTP-sharded" + reference = _full_weights(ref_block, ref_pgs.gtp_remat) + del ref_block + + # ---------------- GTP_remat on, weights allocated pre-sharded -------------------------- + gtp_pgs = _init_mpu(gtp_remat_size=gtp_remat_size) + assert gtp_pgs.gtp_remat.size() == gtp_remat_size, "GTP_remat inactive" + gtp_block = _build_block(config, gtp_pgs) + sharded = [p for p in gtp_block.parameters() if is_gtp_param(p)] + assert sharded, "GTP is not active: the block has no GTP-sharded parameter" + assert any( + getattr(p, "pad_length", 0) > 0 for p in sharded + ), "no GTP-sharded weight needed alignment padding; the padded slice path is not covered" + gathered = _full_weights(gtp_block, gtp_pgs.gtp_remat) + + ps.destroy_model_parallel() + ps.initialize_model_parallel() + + assert set(gathered) == set(reference), "parameter sets differ between the two layouts" + for name, values in gathered.items(): + ref = reference[name] + assert values.shape == ref.shape, f"{name}: shape {values.shape} != reference {ref.shape}" + torch.testing.assert_close( + values, + ref, + atol=0.0, + rtol=0.0, + msg=lambda m, name=name: ( + f"{name}: CPU init at GTP_remat={gtp_remat_size} differs from GTP_remat=1\n{m}" + ), + ) + + +def _worker_high_precision_init_val(rank, world_size, port, gtp_remat_size): + """The optimizer's master-weight snapshot must hold exactly the CPU-initialized values. + + Compared against a BF16 build rather than against the dequantized FP8 weight: both come from + the same master weight rounded to BF16, so the match is exact. Dequantizing would only allow an + assertion to within MXFP8 quantization error, which is loose enough to hide a partly-stale + snapshot. + """ + bf16_pgs = _init_mpu(gtp_remat_size=gtp_remat_size) + bf16_block = _build_block(_make_config(fp8_param=False), bf16_pgs) + bf16_weights = {n: p.data.detach().float().cpu() for n, p in bf16_block.named_parameters()} + del bf16_block + + fp8_pgs = _init_mpu(gtp_remat_size=gtp_remat_size) + block = _build_block(_make_config(fp8_param=True), fp8_pgs) + + checked = 0 + for name, param in block.named_parameters(): + getter = getattr(param, "get_high_precision_init_val", None) + if getter is None or getter() is None: + continue + checked += 1 + torch.testing.assert_close( + getter().float().cpu(), + bf16_weights[name], + atol=0.0, + rtol=0.0, + msg=lambda m, name=name: f"{name}: stale high-precision init snapshot\n{m}", + ) + assert checked > 0, "no quantized primary weights found; fp8_param did not take effect" + + from megatron.core import parallel_state as ps + + ps.destroy_model_parallel() + ps.initialize_model_parallel() + + +class TestGTPCpuInitialization: + @pytest.mark.parametrize("fp8_param", [False, True]) + @pytest.mark.parametrize("gtp_remat_size", GTP_SIZES) + def test_cpu_init_is_gtp_invariant(self, gtp_remat_size, fp8_param): + _requires_multi_gpu(WORLD) + if fp8_param: + _requires_mxfp8() + _run_distributed(_worker_cpu_init_is_gtp_invariant, WORLD, fp8_param, gtp_remat_size) + + @pytest.mark.parametrize("gtp_remat_size", [1, *GTP_SIZES]) + def test_high_precision_init_val_refreshed(self, gtp_remat_size): + """Also covered at GTP=1: the stale snapshot is a plain cpu-init + fp8_param bug, and + nothing about it needs GTP.""" + _requires_multi_gpu(WORLD) + _requires_mxfp8() + _run_distributed(_worker_high_precision_init_val, WORLD, gtp_remat_size)