dist_ckpt: inference-only streaming checkpoint dequantization for --fp8-param-gather / --fp4-param-gather loads (#7405)

Signed-off-by: asolergi-nv <asolergibert@nvidia.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

Author: asolergi-nv

megatron/core/dist_checkpointing/exchange_utils.py

@@ -18,7 +18,7 @@
 from .core import CheckpointingException
 from .dict_utils import nested_values
 from .mapping import ShardedStateDict, ShardedTensor, is_main_replica
-from .utils import _sharded_tensor_shard_id, _ShardId, debug_time
+from .utils import _quantized_storage_tensors, _sharded_tensor_shard_id, _ShardId, debug_time
 
 # TODO: remove TE references once the TE bug is fixed
 # Check if Transformer Engine has Float8Tensor class
@@ -764,12 +764,147 @@ def exchange_loaded_tensors_broadcast(
     return all_loaded_tensors
 
 
+_QUANTIZED_STORAGE_PART_ALIGNMENT = 16
+
+
+def _aligned_nbytes(nbytes: int) -> int:
+    return -(-nbytes // _QUANTIZED_STORAGE_PART_ALIGNMENT) * _QUANTIZED_STORAGE_PART_ALIGNMENT
+
+
+def _quantized_storage_nbytes(tensor: torch.Tensor) -> int:
+    """Size of the packed quantized storage of a TE quantized tensor (see `_pack_...`)."""
+    return sum(
+        _aligned_nbytes(part.numel() * part.element_size())
+        for part in _quantized_storage_tensors(tensor)
+    )
+
+
+def _pack_quantized_storage(tensor: torch.Tensor) -> torch.Tensor:
+    """Packs the raw storage of a TE quantized tensor (codes and scales) into one byte buffer.
+
+    Exchanging the raw storage instead of the dequantized values keeps the loading rank's
+    encoding bit for bit on the receivers (re-quantizing dequantized values can pick different
+    block scales) and halves the exchanged bytes. Parts are 16-byte aligned so that every part
+    can be viewed in place with its own dtype.
+    """
+    buf = torch.empty(_quantized_storage_nbytes(tensor), dtype=torch.uint8, device=tensor.device)
+    offset = 0
+    for part in _quantized_storage_tensors(tensor):
+        nbytes = part.numel() * part.element_size()
+        buf[offset : offset + nbytes].view(part.dtype).view(part.shape).copy_(part)
+        offset += _aligned_nbytes(nbytes)
+    return buf
+
+
+def _unpack_quantized_storage(buf: torch.Tensor, tensor: torch.Tensor) -> None:
+    """Inverse of `_pack_quantized_storage`, writing into the storage of `tensor`."""
+    offset = 0
+    for part in _quantized_storage_tensors(tensor):
+        nbytes = part.numel() * part.element_size()
+        part.copy_(buf[offset : offset + nbytes].view(part.dtype).view(part.shape))
+        offset += _aligned_nbytes(nbytes)
+    if offset != buf.numel():
+        raise CheckpointingException(
+            f'Packed quantized storage ({buf.numel()} bytes) does not match the destination'
+            f' ({offset} bytes)'
+        )
+
+
+def exchange_loaded_tensors_broadcast_quantized(
+    loaded_tensors: Dict[_ShardId, torch.Tensor],
+    unloaded_shards: Dict[_ShardId, ShardedTensor],
+    shard_distribution: ShardDistribution,
+    parallelization_group: Optional[torch.distributed.ProcessGroup] = None,
+) -> Dict[_ShardId, torch.Tensor]:
+    """`exchange_loaded_tensors_broadcast` for state dicts with quantized destinations.
+
+    Used by the streaming dequantize load, which loads quantized destinations in place on the
+    loading rank. Such tensors are broadcast as their packed raw storage (see
+    `_pack_quantized_storage`) and unpacked straight into the receivers' destinations, so that
+    every rank ends up with the loading rank's encoding and no high-precision copy is created.
+    Other tensors are exchanged as by `exchange_loaded_tensors_broadcast`.
+
+    Args: see `exchange_loaded_tensors_broadcast`.
+    """
+    from ..fp8_utils import is_float8tensor  # Avoid circular import
+
+    main_rank_for_shard, _, shard_to_metadata, all_ranks_for_shard = shard_distribution
+    local_rank = torch.distributed.get_rank(group=parallelization_group)
+
+    # Ranks that do not need a shard still take part in its broadcast and must know the size of
+    # its packed storage, which only the loading rank knows: gather the sizes once up front.
+    local_nbytes = {
+        shard_id: _quantized_storage_nbytes(tensor)
+        for shard_id, tensor in loaded_tensors.items()
+        if is_float8tensor(tensor)
+    }
+    all_nbytes = [None] * torch.distributed.get_world_size(group=parallelization_group)
+    torch.distributed.all_gather_object(all_nbytes, local_nbytes, group=parallelization_group)
+    quantized_nbytes = {k: v for rank_nbytes in all_nbytes for k, v in rank_nbytes.items()}
+
+    if not quantized_nbytes:
+        return exchange_loaded_tensors_broadcast(
+            loaded_tensors, unloaded_shards, shard_distribution, parallelization_group
+        )
+
+    plain_loaded_tensors = {k: v for k, v in loaded_tensors.items() if k not in quantized_nbytes}
+    plain_distribution = ShardDistribution(
+        {k: v for k, v in main_rank_for_shard.items() if k not in quantized_nbytes},
+        shard_distribution.shards_in_this_group,
+        shard_to_metadata,
+        all_ranks_for_shard,
+    )
+    all_loaded_tensors = exchange_loaded_tensors_broadcast(
+        plain_loaded_tensors, unloaded_shards, plain_distribution, parallelization_group
+    )
+
+    for shard_id, rank in main_rank_for_shard.items():
+        if shard_id not in quantized_nbytes:
+            continue
+        nbytes = quantized_nbytes[shard_id]
+        if rank == local_rank:
+            all_loaded_tensors[shard_id] = loaded_tensors[shard_id]
+        if len(all_ranks_for_shard[shard_id]) == 1:
+            # Only the loading rank needs this tensor, no exchange needed.
+            continue
+        if rank == local_rank:
+            quantized = loaded_tensors[shard_id]
+            buf = _pack_quantized_storage(quantized)
+        else:
+            quantized = unloaded_shards[shard_id].data if shard_id in unloaded_shards else None
+            if quantized is not None and (
+                not is_float8tensor(quantized) or _quantized_storage_nbytes(quantized) != nbytes
+            ):
+                raise CheckpointingException(
+                    f'Shard {shard_id} was loaded as a quantized tensor on rank {rank} but its'
+                    f' destination on rank {local_rank} does not have the same quantized storage'
+                )
+            buf = torch.empty(nbytes, dtype=torch.uint8, device='cuda')
+        global_src_rank = (
+            rank
+            if parallelization_group is None
+            else torch.distributed.get_global_rank(parallelization_group, rank)
+        )
+        torch.distributed.broadcast(buf, src=global_src_rank, group=parallelization_group)
+        if rank != local_rank and quantized is not None:
+            with torch.no_grad():
+                _unpack_quantized_storage(buf, quantized)
+            # Hand back a view rather than the destination object itself: `load_state_dict`
+            # copies the loaded tensor into the param, and TE's generic quantized `copy_`
+            # cannot copy a tensor onto itself (it detaches the source storage first).
+            all_loaded_tensors[shard_id] = quantized.view(quantized.shape)
+        del buf
+
+    return all_loaded_tensors
+
+
 def exchange_by_distribution(
     loaded_tensors: Dict[_ShardId, torch.Tensor],
     unloaded_shards: Dict[_ShardId, ShardedTensor],
     shard_distribution: ShardDistribution,
     parallelization_group: Optional[torch.distributed.ProcessGroup] = None,
     exchange_algo="broadcast",
+    exchange_quantized_storage: bool = False,
 ) -> Dict[_ShardId, torch.Tensor]:
     """Exchange tensors loaded by different ranks using the specified exchange_algo.
 
@@ -783,6 +918,10 @@ def exchange_by_distribution(
             distribution. Tensors will be exchanged within this group
         exchange_algo (str): The algorithm used for performing exchanges.
             Defaults to 'broadcast'.
+        exchange_quantized_storage (bool): with the 'broadcast' algorithm, exchange quantized
+            tensors as their raw storage and unpack them into the receivers' quantized
+            destinations (streaming dequantize load), see
+            `exchange_loaded_tensors_broadcast_quantized`.
 
     Returns:
         Dict[_ShardId, torch.Tensor]: dictionary mapping shard ids to tensors
@@ -796,7 +935,11 @@ def exchange_by_distribution(
     elif exchange_algo == "gather_rounds":
         exchange_fn = exchange_loaded_tensors_gather_rounds
     elif exchange_algo == "broadcast":
-        exchange_fn = exchange_loaded_tensors_broadcast
+        exchange_fn = (
+            exchange_loaded_tensors_broadcast_quantized
+            if exchange_quantized_storage
+            else exchange_loaded_tensors_broadcast
+        )
     else:
         raise NotImplementedError(f"Unrecognized gather algorithm: {exchange_algo}")
     return exchange_fn(loaded_tensors, unloaded_shards, shard_distribution, parallelization_group)

megatron/core/dist_checkpointing/serialization.py

@@ -35,7 +35,11 @@
     TorchDistSaveShardedStrategy,
     _get_filesystem_reader,
 )
-from .utils import extract_sharded_base, force_all_tensors_to_non_fp8
+from .utils import (
+    extract_sharded_base,
+    force_all_tensors_to_non_fp8,
+    prepare_quantized_tensors_for_streaming_load,
+)
 from .validation import (
     StrictHandling,
     determine_global_metadata,
@@ -137,7 +141,19 @@ def load(
     #      params with a high-precision state dict;
     #   2. When using delayed scaling, this loading process writes an extra value into the global
     #      amax_history buffer of Transformer Engine, which is undesirable.
-    force_all_tensors_to_non_fp8(sharded_state_dict)
+    #
+    # Dequantizing everything up front allocates a high-precision copy of every quantized param
+    # at once. When the strategy streams the dequantization (`stream_ckpt_dequant`, a load path
+    # for inference), quantized tensors stay in the state dict and the load planner quantizes
+    # each destination in place from a per-tensor high-precision scratch (see
+    # `MCoreLoadPlanner.resolve_tensor`). Case 1 above is then unsupported by construction,
+    # delayed scaling (case 2: its scale is only known once `load_state_dict` restores the fp8
+    # metadata) is rejected, and the few tensors that cannot be quantized in place (see
+    # `prepare_quantized_tensors_for_streaming_load`) are still dequantized up front.
+    if getattr(sharded_strategy, "stream_ckpt_dequant", False):
+        prepare_quantized_tensors_for_streaming_load(sharded_state_dict)
+    else:
+        force_all_tensors_to_non_fp8(sharded_state_dict)
 
     sharded_state_dict, nonpersistent_state_dict, sh_ten_factories = load_preprocess(
         sharded_state_dict

megatron/core/dist_checkpointing/strategies/fully_parallel.py

@@ -218,6 +218,15 @@ def __init__(
         self.cached_distribution: Optional[ShardDistribution] = None
         self.cached_global_metadata: Optional[Metadata] = None
 
+    @property
+    def stream_ckpt_dequant(self) -> bool:
+        """Forwards the streaming dequantize flag of the wrapped strategy (see
+        `TorchDistLoadShardedStrategy`), so that `serialization.load` keeps the quantized
+        destinations in the state dict. Only the `broadcast` exchange supports it: it replicates
+        the raw quantized storage of the loading rank into the receivers' destinations (see
+        `exchange_loaded_tensors_broadcast_quantized`)."""
+        return getattr(self.base_strategy, "stream_ckpt_dequant", False)
+
     @debug_time("FullyParallelLoadStrategyWrapper.load", logger)
     def load(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path) -> StateDict:
         """Distributes the load and calls underlying strategy only for parts of the state dict.
@@ -255,6 +264,12 @@ def load(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path) -> St
         if get_pg_size(self.parallelization_group) <= 1:
             return self.base_strategy.load(sharded_state_dict, checkpoint_dir)
 
+        if self.stream_ckpt_dequant and self.exchange_algo != 'broadcast':
+            raise CheckpointingException(
+                'The streaming dequantize load (stream_ckpt_dequant) supports only the'
+                f' `broadcast` exchange algorithm, got `{self.exchange_algo}`.'
+            )
+
         # Step 1 and 2: exchange load metadata and distribute the load
         with debug_time("self.apply_loading_parallelization", logger):
             precomputed_distribution: ShardDistribution | None = self.apply_loading_parallelization(
@@ -326,6 +341,7 @@ def load(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path) -> St
                 precomputed_distribution,
                 self.parallelization_group,
                 self.exchange_algo,
+                exchange_quantized_storage=self.stream_ckpt_dequant,
             )
             if not set(unloaded_shards.keys()).issubset(all_loaded_tensors.keys()):
                 missing_shards = set(unloaded_shards.keys()) - all_loaded_tensors.keys()