[AutoTP] Complete uneven sharding and universal checkpoint support

deepspeed/checkpoint/constants.py

@@ -58,9 +58,10 @@
 UNIVERSAL_CHECKPOINT_INFO = 'universal_checkpoint_info'
 UNIVERSAL_CHECKPOINT_VERSION_KEY = 'universal_checkpoint_version'
 # Reserve version 0.1  for the hardcoded logic used in BLOOM-176B training
-UNIVERSAL_CHECKPOINT_VERSION_VALUE = 0.3
+UNIVERSAL_CHECKPOINT_VERSION_VALUE = 0.4
 # Attribute name used to store AutoTP universal-checkpoint metadata on torch Parameters.
 DS_AUTOTP_UC_META = "ds_autotp_universal_checkpoint_meta"
+AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS = "autotp_unsupported_parameter_patterns"
 
 # Vocabulary padding
 VOCAB_TENSOR = 'vocab_tensor'
@@ -89,6 +90,11 @@
 TP_REPLICATED_PARAMETER_PATTERNS = 'tp_replicated_parameter_patterns'
 PARAMETER_WITH_2_SUB_PARAMS_CAT_DIM_0 = 'parameter_with_2_sub_params_cat_dim_0'
 PARAMETER_WITH_SUB_PARAMS = 'parameter_with_sub_params'
+# Per-rank width of every sub-parameter, keyed by the same pattern used in
+# PARAMETER_WITH_SUB_PARAMS. Kept as a separate top-level key so that converters predating
+# uneven sub-parameter support simply do not see it, instead of failing to build a
+# SubparamShape from an unexpected field.
+SUB_PARAM_SHARD_WIDTHS = 'sub_param_shard_widths'  # UCP version 0.4
 SUB_PARAMS_SHAPE = 'sub_params_shape'
 
 #########################################

deepspeed/checkpoint/universal_checkpoint.py

@@ -10,7 +10,7 @@
 from typing import List, Tuple, Union
 from dataclasses import dataclass
 from .constants import (FP32_WEIGHT_KEY, PARAM, VOCAB_TENSOR, CAT_DIM, PARAM_N_SUB_PARAMS, SUB_PARAM_SHAPE,
-                        EP_IS_EXPERT_PARAM, EP_NUM_EXPERTS, DS_AUTOTP_UC_META)
+                        EP_IS_EXPERT_PARAM, EP_NUM_EXPERTS, DS_AUTOTP_UC_META, UNIVERSAL_CHECKPOINT_VERSION_KEY)
 
 
 @dataclass
@@ -31,6 +31,46 @@ def _get_param_uc_restore_meta(param):
     return getattr(param, DS_AUTOTP_UC_META, None)
 
 
+def _narrow_sub_params(full_view, partition_dim, sub_dim_sizes, shard_widths, tp_rank, tp_world_size, uc_version):
+    """Take this rank's piece of every sub-parameter and concatenate them back together.
+
+    ``shard_widths`` gives the per-rank width of each sub-parameter, so a fused attention
+    weight is cut on key/value head boundaries. Metadata written before UCP version 0.4 carries
+    no widths and can only have been split evenly.
+    """
+    assert sum(sub_dim_sizes) == full_view.shape[partition_dim], (
+        f"Sub-parameter sizes {list(sub_dim_sizes)} sum to {sum(sub_dim_sizes)}, but dimension {partition_dim} "
+        f"of the parameter is {full_view.shape[partition_dim]}. The metadata describes a different parameter.")
+
+    if uc_version < 0.4:
+        # An even split is the only layout recoverable without recorded widths. Assuming it for
+        # an uneven sub-parameter shifts every offset and silently drops the trailing elements,
+        # so refuse rather than restore wrong weights.
+        uneven = [size for size in sub_dim_sizes if size % tp_world_size != 0]
+        assert not uneven, (
+            f"Sub-parameter sizes {uneven} are not divisible by tp_world_size {tp_world_size}, and this "
+            "metadata records no sub_param_shard_widths, so its uneven layout cannot be reconstructed.")
+        shard_widths = [[size // tp_world_size] * tp_world_size for size in sub_dim_sizes]
+
+    assert len(shard_widths) == len(sub_dim_sizes), (
+        f"Got {len(shard_widths)} shard width entries for {len(sub_dim_sizes)} sub-parameters.")
+
+    offset = 0
+    merged_chunks = []
+    for sub_dim_size, widths in zip(sub_dim_sizes, shard_widths):
+        assert len(widths) == tp_world_size, (
+            f"Sub-parameter of size {sub_dim_size} has {len(widths)} shard widths, expected one per tp rank "
+            f"({tp_world_size}).")
+        assert sum(widths) == sub_dim_size, (
+            f"Sub-parameter shard widths {list(widths)} sum to {sum(widths)}, expected {sub_dim_size}.")
+        start = offset + sum(widths[:tp_rank])
+        merged_chunks.append(full_view.narrow(partition_dim, start, widths[tp_rank]))
+        offset += sub_dim_size
+
+    slice_tensor = torch.cat(merged_chunks, dim=partition_dim)
+    return slice_tensor.flatten()
+
+
 def _resolve_autotp_partition(current_param, ckpt_dict, full_hp_param, tp_rank, tp_world_size):
     meta = _get_param_uc_restore_meta(current_param)
     if not meta:
@@ -40,7 +80,17 @@ 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')
+    sub_param_shard_widths = meta.get('sub_param_shard_widths')
+    partition_sizes = meta.get('partition_sizes')
     replicated = meta.get('replicated', False)
+    uc_version = meta.get(UNIVERSAL_CHECKPOINT_VERSION_KEY, 0.0)
+
+    # The layer could not describe how it split this parameter, so conversion refuses it. The
+    # generic paths below would reassemble it as contiguous rank-ordered slices, which is not
+    # how a fused layout was cut, so restoring here would silently corrupt the weight.
+    unsupported_reason = meta.get('conversion', {}).get('unsupported_reason')
+    if unsupported_reason:
+        raise RuntimeError(f"Cannot restore a universal checkpoint into this AutoTP parameter: {unsupported_reason}")
 
     if replicated:
         assert partition_dim is None
@@ -55,43 +105,43 @@ def _resolve_autotp_partition(current_param, ckpt_dict, full_hp_param, tp_rank,
 
     full_view = full_hp_param.view(logical_shape)
 
-    if sub_param_shape is not None:
+    # sub_param_sizes holds the resolved physical width of each sub-parameter, whereas
+    # sub_param_shape may be a logical view spec such as (3, -1) whose partition_dim entry is
+    # the sub-parameter *count*. Reading that count as a width restores only a fraction of the
+    # parameter, so only fall back to sub_param_shape for older metadata that lacks the sizes.
+    sub_dim_sizes = None
+    if sub_param_sizes is not None:
+        sub_dim_sizes = sub_param_sizes
+    elif sub_param_shape is not None:
         if hasattr(sub_param_shape, "shape") and hasattr(sub_param_shape, "partition_dim"):
             shape_spec = sub_param_shape.shape
             partition_dim = sub_param_shape.partition_dim
         else:
             shape_spec = sub_param_shape
-
         sub_dim_sizes = shape_spec[partition_dim]
-        if not isinstance(sub_dim_sizes, tuple):
-            sub_dim_sizes = (sub_dim_sizes, )
 
-        offset = 0
-        merged_chunks = []
-        for sub_dim_size in sub_dim_sizes:
-            sub_slice = full_view.narrow(partition_dim, offset, sub_dim_size) \
-                                .chunk(tp_world_size, dim=partition_dim)[tp_rank]
-            merged_chunks.append(sub_slice)
-            offset += sub_dim_size
+    if sub_dim_sizes is not None:
+        if not isinstance(sub_dim_sizes, (tuple, list)):
+            sub_dim_sizes = (sub_dim_sizes, )
 
-        slice_tensor = torch.cat(merged_chunks, dim=partition_dim)
-        return slice_tensor.flatten()
+        return _narrow_sub_params(full_view, partition_dim, sub_dim_sizes, sub_param_shard_widths, tp_rank,
+                                  tp_world_size, uc_version)
 
-    if sub_param_sizes is not None:
-        if not isinstance(sub_param_sizes, (tuple, list)):
-            sub_param_sizes = (sub_param_sizes, )
-
-        offset = 0
-        merged_chunks = []
-        for sub_dim_size in sub_param_sizes:
-            sub_slice = full_view.narrow(partition_dim, offset, sub_dim_size) \
-                                .chunk(tp_world_size, dim=partition_dim)[tp_rank]
-            merged_chunks.append(sub_slice)
-            offset += sub_dim_size
-
-        slice_tensor = torch.cat(merged_chunks, dim=partition_dim)
+    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()
 
+    # torch.chunk sizes every block as ceil(size / tp_world_size) and shrinks the last one,
+    # while AutoTP hands one extra element to each of the first `size % tp_world_size` ranks.
+    # The two disagree once the split is uneven (10/4 -> chunk [3, 3, 3, 1] vs AutoTP
+    # [3, 3, 2, 2]), and chunk can even return fewer blocks than there are ranks.
+    partition_dim_size = full_view.shape[partition_dim]
+    assert partition_dim_size % tp_world_size == 0, (
+        f"Dimension {partition_dim} of size {partition_dim_size} is not divisible by tp_world_size "
+        f"{tp_world_size}, and this checkpoint records no partition_sizes, so its uneven layout cannot "
+        "be reconstructed.")
     slice_tensor = full_view.chunk(tp_world_size, dim=partition_dim)[tp_rank]
     return slice_tensor.flatten()
 
@@ -148,15 +198,21 @@ def load_hp_checkpoint_state(self, folder, tp_rank, tp_world_size, ep_rank=0, ep
         # implementation b.
         # this version requires no additional data passed from the client
         # if the shapes already match it must be slices that were averaged - so we just hack around those
-        if full_hp_param.shape == self.shape:
+        # AutoTP restore metadata already pins down this rank's slice of the universal tensor,
+        # including layouts where one rank owns the whole thing (e.g. a single KV head under MQA).
+        # Collapsing the TP topology here would contradict the recorded per-rank widths.
+        has_autotp_meta = _get_param_uc_restore_meta(self) is not None
+        if full_hp_param.shape == self.shape and not has_autotp_meta:
             tp_rank = 0
             tp_world_size = 1
 
         # special case for word_embeddings weights which get padded differently depending on TP degree.
         # 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.
+        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:

deepspeed/compile/custom_ops/tp_collectives.py

@@ -47,30 +47,46 @@ def reduce_from_tp_region_fake(input: torch.Tensor):
 
 
 @torch.library.custom_op("autotp::gather_from_tp_region", mutates_args=())
-def gather_from_tp_region(input: torch.Tensor) -> torch.Tensor:
-    """All-gather the last dimension in the forward pass, take this rank's slice in the backward.
+def gather_from_tp_region(input: torch.Tensor, partition_sizes: list[int]) -> torch.Tensor:
+    """All-gather the last dimension using the frozen shard widths.
 
     Inserted after a column-parallel matmul whose layer asks for gather_output, so that
-    every rank leaves the layer holding the full output rather than its own shard.
+    every rank leaves the layer holding the full output rather than its own shard. Unlike the
+    even-width fast path, this takes an explicit `partition_sizes` list so uneven TP shards are
+    still reconstructed correctly.
     """
     group = get_tp_group()
     world_size = dist.get_world_size(group=group)
     if world_size == 1:
         return input.clone()
 
-    local_shard = input.contiguous()
+    if len(partition_sizes) != world_size:
+        raise ValueError(f"partition_sizes={partition_sizes} does not match TP world size {world_size}")
+
+    local_rank = dist.get_rank(group=group)
+    local_size = partition_sizes[local_rank]
+    if input.shape[-1] != local_size:
+        raise ValueError(
+            f"Rank {local_rank} produced width {input.shape[-1]}, but partition_sizes expects {local_size}")
+
+    max_partition_size = max(partition_sizes)
+    if local_size == max_partition_size:
+        local_shard = input.contiguous()
+    else:
+        local_shard = input.new_zeros((*input.shape[:-1], max_partition_size))
+        local_shard[..., :local_size].copy_(input)
+
     flat_gathered = torch.empty((world_size * local_shard.shape[0], *local_shard.shape[1:]),
                                 dtype=local_shard.dtype,
                                 device=local_shard.device)
     dist.all_gather_into_tensor(flat_gathered, local_shard, group=group)
     shards = flat_gathered.view(world_size, *local_shard.shape)
-    return torch.cat(shards.unbind(0), dim=-1)
+    return torch.cat([shards[i].narrow(-1, 0, size) for i, size in enumerate(partition_sizes)], dim=-1)
 
 
 @torch.library.register_fake("autotp::gather_from_tp_region")
-def gather_from_tp_region_fake(input: torch.Tensor):
-    world_size = dist.get_world_size(group=get_tp_group())
-    return input.new_empty((*input.shape[:-1], input.shape[-1] * world_size))
+def gather_from_tp_region_fake(input: torch.Tensor, partition_sizes: list[int]):
+    return input.new_empty((*input.shape[:-1], sum(partition_sizes)))
 
 
 def _copy_to_tp_region_backward(ctx, grad):
@@ -83,14 +99,22 @@ def _reduce_from_tp_region_backward(ctx, grad):
     return grad
 
 
+def _gather_from_tp_region_setup(ctx, inputs, output):
+    _, partition_sizes = inputs
+    ctx.partition_sizes = tuple(partition_sizes)
+
+
 def _gather_from_tp_region_backward(ctx, grad):
     group = get_tp_group()
     world_size = dist.get_world_size(group=group)
     if world_size == 1:
-        return grad
-    shard_width = grad.shape[-1] // world_size
-    shard_start = dist.get_rank(group=group) * shard_width
-    return grad.narrow(-1, shard_start, shard_width).contiguous()
+        return (grad, None)
+
+    partition_sizes = ctx.partition_sizes
+    rank = dist.get_rank(group=group)
+    shard_start = sum(partition_sizes[:rank])
+    shard_width = partition_sizes[rank]
+    return grad.narrow(-1, shard_start, shard_width).contiguous(), None
 
 
 def _setup_context_without_saved_tensors(ctx, inputs, output):
@@ -105,4 +129,4 @@ def _setup_context_without_saved_tensors(ctx, inputs, output):
                                 setup_context=_setup_context_without_saved_tensors)
 torch.library.register_autograd("autotp::gather_from_tp_region",
                                 _gather_from_tp_region_backward,
-                                setup_context=_setup_context_without_saved_tensors)
+                                setup_context=_gather_from_tp_region_setup)