# mypy: allow-untyped-defs # Copyright (c) Meta Platforms, Inc. and affiliates import logging from typing import Any, Dict, List, Optional, Tuple import torch from torch.fx.node import map_aggregate from torch.utils._pytree import tree_flatten, tree_unflatten __all__ = [ "TensorChunkSpec", "split_args_kwargs_into_chunks", "merge_chunks", ] logger = logging.getLogger(__name__) """ _debug_mask_minibatches specifies to send masked versions of the mini-batch through instead of micro-batch slices--this can be used for more stable numerical testing (see [A Note About Correctness Testing]) """ _debug_mask_minibatches = False class _CustomReducer: """ Custom reducer class that can be used to specify a custom operation that reduces losses of multiple microbatches into one value. Example: >>> # xdoctest: +SKIP >>> sum_reducer = _CustomReducer( >>> torch.tensor(0.0), >>> lambda a, b: a + b >>> ) """ def __init__(self, init_value, reduce_fn): self.init_value = init_value self.reduce_fn = reduce_fn class _LossReducer(_CustomReducer): pass sum_reducer = _LossReducer(torch.tensor(0.0), lambda a, b: a + b) # Default chunking dimension is 0. This is used for the case where the user did # not specify a chunking dimension. DEFAULT_CHUNK_DIM = 0 class TensorChunkSpec: """ Class used to specify chunking of inputs """ def __init__(self, split_dim): self.split_dim = split_dim split_dim: int def __repr__(self): return ( f"{self.__class__.__module__}.{self.__class__.__name__}({self.split_dim})" ) def __str__(self): return f"TensorChunkSpec({self.split_dim})" @staticmethod def from_tuple( chunk_dims: Tuple[int, ...], ): """ A helper for creating a tuple of `TensorChunkSpec` from a tuple of chunk dimensions (int's). Example: >>> # xdoctest: +SKIP >>> # There are three positional arguments to the model, and >>> # we are chunking them along dimension 0, 0 and 1, respectively >>> args_chunk_spec = TensorChunkSpec.from_tuple((0, 0, 1)) """ args_chunk_spec = map_aggregate( chunk_dims, lambda dim: TensorChunkSpec(dim), # type: ignore[arg-type,return-value] ) return args_chunk_spec @staticmethod def from_dict( chunk_dims: Dict[str, int], ): """ A helper for creating a dictionary of `TensorChunkSpec` from a dictionary of chunk dimensions (int's). Example: >>> # xdoctest: +SKIP >>> # Chunk dimension 0 for the "id" argument, 1 for the "mask" argument >>> kwargs_chunk_spec = TensorChunkSpec.from_dict({"id": 0, "mask": 1}) """ kwargs_chunk_spec = map_aggregate( chunk_dims, lambda dim: TensorChunkSpec(dim), # type: ignore[arg-type,return-value] ) return kwargs_chunk_spec # Class used to specify replication of inputs class _Replicate: pass def _shard_dict_of_args( args_dict, args_chunk_spec, num_chunks, ): """ Given a dictionary of args, and a dictionary of chunking specs, shard the args according to the chunking specs. Args: args_dict: Dictionary of args args_chunk_spec: Dictionary of chunking specs num_chunks: Number of chunks to shard the args into Returns: args_split: List of sharded args """ # Stage 1+2: flatten and shard/replicate # args_sharded_replicated : [num args, num flat values, num chunks] args_sharded_replicated = {} arg_specs = [] real_num_chunks = num_chunks first_tensor = True assert len(args_dict) == len( args_chunk_spec ), f"args_dict.keys() = {list(args_dict.keys())} args_chunk_spec.keys() = {list(args_chunk_spec.keys())}" for arg_key, arg in args_dict.items(): flat, spec = tree_flatten(arg) arg_specs.append(spec) chunk_spec = args_chunk_spec[arg_key] assert chunk_spec is not None # Should have been set by caller chunk_spec_flat, _ = tree_flatten(chunk_spec) if len(flat) != len(chunk_spec_flat): raise ValueError( f"Argument value {arg} did not have the same number of " f"values as as chunk spec {chunk_spec}" ) sharded_arg_flat = [] for v, chunk_v in zip(flat, chunk_spec_flat): if chunk_v is _Replicate or not isinstance(v, torch.Tensor): sharded_arg_flat.append([v] * real_num_chunks) elif isinstance(chunk_v, TensorChunkSpec): # TODO: check type of v. If it's a tensor, use chunk (or debug mask). # If it's a collection type, split it as you would expect. Otherwise, # Throw an error assert isinstance(v, torch.Tensor), f"{v} is not a tensor" v_split_dim_size = v.size(chunk_v.split_dim) if v_split_dim_size < real_num_chunks: if first_tensor: # We can only adjust number of chunks when we hit this # issue at the first tensor encountered logger.warning( f"Tensor size on chunking dimension is {v_split_dim_size}, " # noqa: G004 f"downsizing the number of chunks from {num_chunks} to {v_split_dim_size}." ) real_num_chunks = v_split_dim_size else: raise RuntimeError( f"Arg {arg_key} on chunking dimension has a size of {v_split_dim_size}, " f"smaller than the number of chunks {num_chunks}. " "PiPPy cannot reduce the number of chunks because " "other arguments have bigger chunk-dimension sizes. " "Please adjust your num_chunks setting." ) chunk_tensors = torch.tensor_split( v, real_num_chunks, chunk_v.split_dim ) if _debug_mask_minibatches: expanded_chunks = [] split_dim_idx = 0 for chunk_tensor in chunk_tensors: new_val = torch.zeros_like(v) upper_idx = split_dim_idx + chunk_tensor.size(chunk_v.split_dim) slice_indices = [slice(None, None, None)] * new_val.ndim slice_indices[chunk_v.split_dim] = slice( split_dim_idx, upper_idx ) new_val[slice_indices] = chunk_tensor expanded_chunks.append(new_val) split_dim_idx += chunk_tensor.size(chunk_v.split_dim) sharded_arg_flat.append(expanded_chunks) else: sharded_arg_flat.append(chunk_tensors) # type: ignore[arg-type] first_tensor = False else: raise TypeError(f"Unrecognized chunk spec: {chunk_v}") args_sharded_replicated[arg_key] = sharded_arg_flat # chunks_flat : [num chunks, num args, num flat values] chunks_flat = [] for chunk_idx in range(real_num_chunks): chunk_args = {} for key, arg in args_sharded_replicated.items(): arg_single_chunk = [] for v_flat in arg: arg_single_chunk.append(v_flat[chunk_idx]) chunk_args[key] = arg_single_chunk chunks_flat.append(chunk_args) # args_split : [num chunks, num args] args_split = [] for chunk in chunks_flat: per_chunk_args = {} assert len(arg_specs) == len(chunk) for (key, arg), arg_spec in zip(chunk.items(), arg_specs): per_chunk_args[key] = tree_unflatten(arg, arg_spec) args_split.append(per_chunk_args) return args_split def split_args_kwargs_into_chunks( args: Tuple[Any, ...], kwargs: Optional[Dict[str, Any]], chunks: int, args_chunk_spec: Optional[Tuple[TensorChunkSpec, ...]] = None, kwargs_chunk_spec: Optional[Dict[str, TensorChunkSpec]] = None, ) -> Tuple[List[Tuple], List[Dict]]: """ Given a sequence of args and kwargs, split them into a number of chunks according to their respective chunking specs. Args: args: Tuple of args kwargs: Dict of kwargs chunks: Number of chunks to split the args and kwargs into args_chunk_spec: chunking specs for args, in same shape as args kwargs_chunk_spec: chunking specs for kwargs, in same shape as kwargs Returns: args_split: List of sharded args kwargs_split: List of sharded kwargs """ # Given `args` and `kwargs`, we want to yield a set of `chunks` args and kwargs such that # the constituent Tensor values have been sharded/replicated according to the `args_chunk_spec` # and `kwargs_chunk_spec` specifications. The steps are as follows: # # 1. Use pytree.tree_flatten to flatten each arg and its spec into nto a 1d array of values. # To use a running example: suppose our inputs look like # # args = ([A, [B, C]], D) args_spec = ([None, [None, TensorChunkSpec]], None) # (kwargs not shown but it's a similar process) # # Then for this step we would end up with # # args = ([A, B, C], D) args_spec = ([None, None, TensorChunkSpec], None) # # 2. Shard or replicate the arguments subject to the policy in the spec. Suppose chunks = 2 # # args = ([[A, A], [B, B], [C_1, C_2]], [D, D]) # # 3. Rotate the nesting order such that chunks are the outer dimension # # args_chunks = [ # ([A, B, C_1], D), # ([A, B, C_2], D), # ] # # 4. Unflatten each chunk according to the spec # # args_chunks = [ # ([A, [B, C_1]], D), # ([A, [B, C_2]], D), # ] # TODO: _debug_mask_minibatches # Handle the case where kwargs is None if kwargs is None: kwargs = {} # If user did not provide args_chunk_spec or kwargs_chunk_spec, we extend # their format and use default chunking along dim 0 if args_chunk_spec is None: args_chunk_spec = (TensorChunkSpec(DEFAULT_CHUNK_DIM),) * len(args) if kwargs_chunk_spec is None: kwargs_chunk_spec = dict.fromkeys(kwargs, TensorChunkSpec(DEFAULT_CHUNK_DIM)) args_split_dict = _shard_dict_of_args( dict(enumerate(args)), dict(enumerate(args_chunk_spec)), chunks, ) real_num_chunks = len(args_split_dict) kwargs_split = _shard_dict_of_args( kwargs, kwargs_chunk_spec, real_num_chunks, ) if len(kwargs_split) < real_num_chunks: # In case kwargs are sharded into less chunks # e.g. when `args` has no tensor, just values real_num_chunks = len(kwargs_split) # Re-shard args args_split_dict = _shard_dict_of_args( dict(enumerate(args)), dict(enumerate(args_chunk_spec)), real_num_chunks, ) if len(args_split_dict) != len(kwargs_split): raise RuntimeError( "args and kwargs are split into different number of chunks: " f"{len(args_split_dict)}, {len(kwargs_split)}" ) args_split = [] for chunk_args in args_split_dict: args_split.append(tuple(chunk_args[i] for i in range(len(chunk_args)))) return args_split, kwargs_split def merge_chunks( chunks: List[Any], chunk_spec, ): """ Given a list of chunks, merge them into a single value according to the chunk spec. Args: chunks: list of chunks chunk_spec: Chunking spec for the chunks Returns: value: Merged value """ # This is essentially the inverse of `split_args_kwargs_into_chunks`, so the # steps are similar to the steps in that function but in reverse. Given the # input values: # # chunks = [ # ([A, [B, C_1]], D), # ([A, [B, C_2]], D), # ] # args_spec = ([None, [None, TensorChunkSpec]], None) # # 1. Flatten the chunks according to the chunk_spec # # chunks_flat = [ # ([A, B, C_1], D), # ([A, B, C_2], D), # ] # # 2. Rotate the nesting order such that chunks are the inner dimension # # value_inner = ([A, B, [C_1, C_2]], D) # # 3. Concatenate sharded arguments # # value_combined = ([A, B, C], D) # # 4. Unflatten the combined args given the spec # # value = ([A, [B, C]], D) # Preliminary: flatten the chunk spec if chunk_spec is not None: spec_flattened, flatten_spec = tree_flatten(chunk_spec) else: # If chunk_spec is not provided, we will merge chunks along the default dimension (0), for all output fields # We obtain the output structure by flattening chunk 0 and generate the chunk_spec chunk0_flat, flatten_spec = tree_flatten(chunks[0]) spec_flattened = [TensorChunkSpec(DEFAULT_CHUNK_DIM)] * len(chunk0_flat) # Stage 1: flatten chunks # chunks_flattened : [num chunks, num args] chunks_flattened = [] for chunk in chunks: chunk_flattened, _ = tree_flatten(chunk) if len(chunk_flattened) != len(spec_flattened): raise ValueError(f"Chunk {chunk} did not match chunk spec {chunk_spec}") chunks_flattened.append(chunk_flattened) # Stage 2 and 3: Rotate nesting order s.t. chunks are inner dimension and # concatenate sharded operands # args_flattened : [num args] args_flattened = [] for arg_idx, arg in enumerate(spec_flattened): if isinstance(arg, TensorChunkSpec): partial_values = [ chunks_flattened[chunk_idx][arg_idx] for chunk_idx in range(len(chunks_flattened)) ] if _debug_mask_minibatches: # Infer size of individual chunks by running `tensor_split` again overall_shape = partial_values[0].shape for val in partial_values[1:]: assert val.shape == overall_shape meta_chunks = torch.tensor_split( torch.empty(*overall_shape, device="meta"), sections=len(partial_values), dim=arg.split_dim, ) values_to_cat = [] chunk_start_idx = 0 assert len(partial_values) == len(meta_chunks) for partial_value, meta_chunk in zip(partial_values, meta_chunks): chunk_end_idx = chunk_start_idx + meta_chunk.size(arg.split_dim) slice_indices = [slice(None, None, None)] * partial_value.ndim slice_indices[arg.split_dim] = slice(chunk_start_idx, chunk_end_idx) sliced = partial_value[slice_indices] values_to_cat.append(sliced) chunk_start_idx = chunk_end_idx else: values_to_cat = partial_values args_flattened.append(torch.cat(values_to_cat, dim=arg.split_dim)) elif isinstance(arg, _CustomReducer): reduced_val = arg.init_value for chunk_idx in range(len(chunks_flattened)): reduced_val = arg.reduce_fn( reduced_val, chunks_flattened[chunk_idx][arg_idx] ) args_flattened.append(reduced_val) else: value = chunks_flattened[0][arg_idx] for chunk_idx in range(1, len(chunks_flattened)): assert chunks_flattened[chunk_idx][arg_idx] == value args_flattened.append(value) # Stage 4: Unflatten combined args return tree_unflatten(args_flattened, flatten_spec)