Skip to content

(feat) gtp+cpu-init support for dense-layers #6986

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
fanshiqing merged 1 commit into from
Sep 1, 2026
Merged

(feat) gtp+cpu-init support for dense-layers #6986

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 51 additions & 0 deletions megatron/core/tensor_parallel/generalized_tensor_parallelism.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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="<unnamed>"):
"""Pad + slice a full-size BF16 weight to this rank's GTP shard.

Expand Down Expand Up @@ -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_<Fp8TensorClass>`` subclasses, keyed by the FP8 base class.
Expand Down
2 changes: 2 additions & 0 deletions megatron/core/tensor_parallel/gtp_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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",
Expand Down
13 changes: 12 additions & 1 deletion megatron/core/tensor_parallel/layers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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)

@fanshiqing fanshiqing Aug 31, 2026 •
edited
Loading

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note: It's a latent bug only triggered when combine --use-cpu-intiliazation + fp8-param-gather + training mode and this bug is orthogonal to GTP entirely.

Image

if return_master_weight:
return master_weight
return None
Expand Down
81 changes: 81 additions & 0 deletions tests/unit_tests/generalized_tensor_parallel/test_gtp_basics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Loading
Loading