Skip to content

[WIP][AutoTP] Complete uneven sharding and universal checkpoint support #8185

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

Open
jinyouzhi wants to merge 12 commits into
base: master
Choose a base branch
from
Open

[WIP][AutoTP] Complete uneven sharding and universal checkpoint support #8185

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
2936b72
Enhance AutoTP Uneven Column Parallel Gathered Output
jinyouzhi Jul 27, 2026
847c4e4
Complete AutoTP uneven sharding for row parallelism and checkpoints
jinyouzhi Jul 27, 2026
8acd22a
Fix tied parameter shape collection in universal checkpoints
jinyouzhi Jul 28, 2026
be4c6d7
make GatherFromTensorParallelRegion partition sizes recomputed locally
jinyouzhi Jul 29, 2026
790dab9
Remove the diagnostics or change as logger.info
jinyouzhi Aug 2, 2026
000c92f
Freeze AutoTP partition sizes at layer construction
jinyouzhi Aug 2, 2026
f3bf28d
Merge remote-tracking branch 'origin/master' into uneven
jinyouzhi Aug 3, 2026
f9bef4f
Merge remote-tracking branch 'origin/master' into uneven
jinyouzhi Aug 3, 2026
de6159f
Add shape consistency guard for merging tied parameters in pipeline r…
jinyouzhi Aug 3, 2026
f0ce0b3
Keep the sub-grain remainder when sharding a dimension
jinyouzhi Aug 3, 2026
4953f60
Cleanup the remaining AutoTP lm_head replacement traces
jinyouzhi Aug 3, 2026
c17e2e6
fix yapf formatting
jinyouzhi Aug 3, 2026
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
10 changes: 9 additions & 1 deletion deepspeed/checkpoint/ds_to_universal.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -733,7 +733,15 @@ def _group_per_tp_shapes(slice_shapes_by_tp, pp_degree, tp_degree):
for tp in range(tp_degree):
tp_dict = {}
for pp in range(pp_degree):
tp_dict.update(slice_shapes_by_tp[pp * tp_degree + tp])
stage_shapes = slice_shapes_by_tp[pp * tp_degree + tp]
for name in tp_dict.keys() & stage_shapes.keys():
# A parameter tied across pipeline stages is stored in every stage's file.
# The replicas must agree, otherwise the tie was partitioned inconsistently
# and the update below would silently keep the stage that was read last.
assert tp_dict[name] == stage_shapes[name], (
f"Pipeline replicas of {name} on tp rank {tp} disagree on shape: "
f"{tp_dict[name]} vs {stage_shapes[name]}.")
tp_dict.update(stage_shapes)
per_tp.append(tp_dict)
all_names = set()
for d in per_tp:
Expand Down
12 changes: 11 additions & 1 deletion deepspeed/checkpoint/universal_checkpoint.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ def _resolve_autotp_partition(current_param, ckpt_dict, full_hp_param, tp_rank,
logical_shape = meta.get('logical_shape')
sub_param_shape = meta.get('sub_param_shape')
sub_param_sizes = meta.get('sub_param_sizes')
partition_sizes = meta.get('partition_sizes')
replicated = meta.get('replicated', False)

if replicated:
Expand Down Expand Up @@ -92,6 +93,12 @@ def _resolve_autotp_partition(current_param, ckpt_dict, full_hp_param, tp_rank,
slice_tensor = torch.cat(merged_chunks, dim=partition_dim)
return slice_tensor.flatten()

if partition_sizes is not None:
shard_offset = sum(partition_sizes[:tp_rank])
shard_size = partition_sizes[tp_rank]
slice_tensor = full_view.narrow(partition_dim, shard_offset, shard_size)
return slice_tensor.flatten()

slice_tensor = full_view.chunk(tp_world_size, dim=partition_dim)[tp_rank]
return slice_tensor.flatten()

Expand Down Expand Up @@ -156,7 +163,10 @@ def load_hp_checkpoint_state(self, folder, tp_rank, tp_world_size, ep_rank=0, ep
# the converter to universal currently strips the original padding completely so the saved
# weight is padding-free and we just need to add new padding depending on the target TP
# degree
is_vocab_tensor = ckpt_dict.get(VOCAB_TENSOR, False) and not is_expert_param
# AutoTP restore metadata already describes the exact (possibly uneven) partition layout,
# so the legacy tp-degree-derived vocab padding must not be applied on top of it.
has_autotp_meta = _get_param_uc_restore_meta(self) is not None
is_vocab_tensor = ckpt_dict.get(VOCAB_TENSOR, False) and not is_expert_param and not has_autotp_meta
if is_vocab_tensor:
# In the absence of data passed from the user wrt new padded vocab specific to tp degree
# we can again derive that data by reverse engineering the target shapes like so:
Expand Down
34 changes: 4 additions & 30 deletions deepspeed/module_inject/auto_tp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
# DeepSpeed Team

# Automatic Tensor Parallelism
import logging
import re

from torch import nn
Expand All @@ -16,7 +17,7 @@
from .fusedqkv_utils import require_tp_fused_qkvw
from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list
from deepspeed.utils import groups
from deepspeed.utils.logging import print_dist
from deepspeed.utils.logging import log_dist
from deepspeed.module_inject.layers import is_autotp_training_mode
from .autotp_config import TPLayerSpec, AutoTPConfig, PartitionType

Expand Down Expand Up @@ -427,17 +428,6 @@ def _replace_with_config(self, child, name):
model_type = self._get_model_type()
spec = self.partition_config.find_matching_spec(param_name, model_type)

if any(part in ("lm_head", "embed_out") for part in name.split('.')):
spec_details = None
if spec is not None:
spec_details = {
"patterns": spec.patterns,
"partition_type": spec.partition_type.value,
"gather_output": spec.gather_output,
}
print_dist(f"AutoTP lm_head spec match: parameter={param_name!r}; matched_spec={spec_details!r}",
ranks=[0])

if spec is None:
# No matching spec found
if self.partition_config.strict_mode:
Expand Down Expand Up @@ -475,21 +465,6 @@ def _create_row_parallel_layer(self, module, spec: TPLayerSpec, name: str):

def _create_column_parallel_layer(self, module, spec: TPLayerSpec, name: str):
"""Create column-parallel layer (AllReduce in backward)."""
if spec.gather_output and self.mp_size is not None and self.mp_size > 1:
output_dim = module.weight.shape[0]
if output_dim % self.mp_size != 0:
if any(part in ("lm_head", "embed_out") for part in name.split('.')):
print_dist(
f"AutoTP: '{name}' uses gather_output with uneven output dim {output_dim} and tp_size="
f"{self.mp_size}; falling back to legacy LmHeadLinearAllreduce for checkpoint-safe "
"consolidation.",
ranks=[0],
)
return LmHeadLinearAllreduce(module, self.mp_group)
raise NotImplementedError(
f"AutoTP gather_output requires output dimension divisible by tp_size. Layer '{name}' has "
f"output dim {output_dim} with tp_size={self.mp_size}.")

if self.conv_linear_layer:
return conv_LinearLayer(module, self.mp_group, name=name, gather_output=spec.gather_output)
# Only use fused-QKV heuristics when no partition_config is provided.
Expand All @@ -506,8 +481,6 @@ def _create_column_parallel_layer(self, module, spec: TPLayerSpec, name: str):
partition_dim=spec.get_partition_dim(),
name=name,
)
if spec.gather_output:
print_dist(f"AutoTP: replacing '{name}' with LinearLayer(gather_output=True)", ranks=[0])
return LinearLayer(module, self.mp_group, name=name, gather_output=spec.gather_output)

def _configure_gathered_column_tie_fallbacks(self):
Expand Down Expand Up @@ -550,10 +523,11 @@ def _configure_gathered_column_tie_fallbacks(self):
continue

self._tied_gathered_column_module_names.update((module_name, tied_embedding_name))
print_dist(
log_dist(
f"AutoTP: '{module_name}.weight' is tied to '{tied_embedding_name}.weight'; leaving both modules "
"replicated because coupled vocabulary-parallel embedding is not supported yet.",
ranks=[0],
level=logging.WARNING,
)

self._gathered_column_tie_fallbacks_configured = True
Expand Down
Loading
Loading