| |
| |
|
|
| from dataclasses import dataclass |
|
|
| import torch |
|
|
|
|
| @dataclass(frozen=True) |
| class InvalidVocabTailMask: |
| """A compact additive mask for the tile-aligned invalid tail of the final vocab shard.""" |
|
|
| mask: torch.Tensor |
| tail_width: int |
| shard_width: int |
| num_vocab_shards: int |
|
|
|
|
| def build_invalid_vocab_mask( |
| vocab_size: int, |
| padded_vocab_size: int, |
| max_batch_size: int, |
| *, |
| dtype: torch.dtype = torch.bfloat16, |
| ) -> torch.Tensor | None: |
| """Build an additive logits mask for LM-head vocabulary padding. |
| |
| LM heads may pad output weights so the sharded matmul has legal tile/device |
| dimensions. Those padded columns produce logits, but they are not real token |
| IDs and must be masked before argmax or top-k sampling. |
| """ |
| if vocab_size < 0: |
| raise ValueError(f"vocab_size must be non-negative, got {vocab_size}") |
| if padded_vocab_size < vocab_size: |
| raise ValueError(f"padded_vocab_size ({padded_vocab_size}) must be >= vocab_size ({vocab_size})") |
| if max_batch_size <= 0: |
| raise ValueError(f"max_batch_size must be positive, got {max_batch_size}") |
| if vocab_size == padded_vocab_size: |
| return None |
|
|
| if not torch.empty((), dtype=dtype).is_floating_point(): |
| raise TypeError(f"dtype must be a floating point torch dtype, got {dtype}") |
|
|
| mask = torch.zeros(1, 1, max_batch_size, padded_vocab_size, dtype=dtype) |
| mask[..., vocab_size:] = torch.finfo(dtype).min |
| return mask |
|
|
|
|
| def _validate_cluster_shape(cluster_shape: tuple[int, int] | list[int]) -> tuple[int, int]: |
| cluster_shape = tuple(cluster_shape) |
| if len(cluster_shape) != 2: |
| raise ValueError(f"cluster_shape must have two dimensions, got {cluster_shape}") |
|
|
| rows, cols = int(cluster_shape[0]), int(cluster_shape[1]) |
| if rows <= 0 or cols <= 0: |
| raise ValueError(f"cluster_shape dimensions must be positive, got {cluster_shape}") |
| return rows, cols |
|
|
|
|
| def get_vocab_num_shards( |
| cluster_shape: tuple[int, int] | list[int], |
| sampling_all_gather_axis: int = 0, |
| ) -> int: |
| """Return how many mesh partitions own contiguous slices of the vocab dimension.""" |
| rows, cols = _validate_cluster_shape(cluster_shape) |
|
|
| if rows == 1 and cols == 1: |
| return 1 |
| if rows == 1: |
| return cols |
| if cols == 1: |
| return rows |
| if sampling_all_gather_axis == 0: |
| return rows |
| if sampling_all_gather_axis == 1: |
| return cols |
| raise ValueError(f"sampling_all_gather_axis must be 0 or 1, got {sampling_all_gather_axis}") |
|
|
|
|
| def build_tail_invalid_vocab_mask( |
| vocab_size: int, |
| padded_vocab_size: int, |
| max_batch_size: int, |
| cluster_shape: tuple[int, int] | list[int], |
| sampling_all_gather_axis: int = 0, |
| *, |
| dtype: torch.dtype = torch.bfloat16, |
| tile_size: int = 32, |
| ) -> InvalidVocabTailMask | None: |
| """Build a compact mask for padding that lives only at the final shard tail. |
| |
| The sampling logits are sharded into equal local vocab widths. For model |
| shapes like Qwen3-32B on T3K, all invalid IDs are a small tile-aligned suffix |
| of the last local shard. In that case callers can mask only the local tail |
| slice instead of adding a full-vocab all-zero mask on every device. |
| |
| Returns ``None`` when the invalid range is not a tile-aligned final-shard |
| suffix; callers should use ``build_invalid_vocab_mask`` as the correctness |
| fallback. |
| """ |
| if vocab_size < 0: |
| raise ValueError(f"vocab_size must be non-negative, got {vocab_size}") |
| if padded_vocab_size < vocab_size: |
| raise ValueError(f"padded_vocab_size ({padded_vocab_size}) must be >= vocab_size ({vocab_size})") |
| if max_batch_size <= 0: |
| raise ValueError(f"max_batch_size must be positive, got {max_batch_size}") |
| if tile_size <= 0: |
| raise ValueError(f"tile_size must be positive, got {tile_size}") |
| if vocab_size == padded_vocab_size: |
| return None |
| if not torch.empty((), dtype=dtype).is_floating_point(): |
| raise TypeError(f"dtype must be a floating point torch dtype, got {dtype}") |
|
|
| num_vocab_shards = get_vocab_num_shards(cluster_shape, sampling_all_gather_axis) |
| if padded_vocab_size % num_vocab_shards != 0: |
| return None |
|
|
| shard_width = padded_vocab_size // num_vocab_shards |
| tail_width = padded_vocab_size - vocab_size |
| if tail_width > shard_width: |
| return None |
| if tail_width % tile_size != 0 or (shard_width - tail_width) % tile_size != 0: |
| return None |
|
|
| mask = torch.zeros(1, 1, max_batch_size, tail_width * num_vocab_shards, dtype=dtype) |
| final_tail_start = tail_width * (num_vocab_shards - 1) |
| mask[..., final_tail_start:] = torch.finfo(dtype).min |
| return InvalidVocabTailMask( |
| mask=mask, |
| tail_width=tail_width, |
| shard_width=shard_width, |
| num_vocab_shards=num_vocab_shards, |
| ) |
|
|
|
|
| def get_vocab_shard_dims( |
| cluster_shape: tuple[int, int] | list[int], |
| sampling_all_gather_axis: int = 0, |
| ) -> tuple[int | None, int | None]: |
| """Return the 2D mesh mapper dims for sharding vocab over the sampling TP axis.""" |
| rows, cols = _validate_cluster_shape(cluster_shape) |
|
|
| if rows == 1 and cols == 1: |
| return (None, None) |
| if rows == 1: |
| return (None, 3) |
| if cols == 1: |
| return (3, None) |
| if sampling_all_gather_axis == 0: |
| return (3, None) |
| if sampling_all_gather_axis == 1: |
| return (None, 3) |
| raise ValueError(f"sampling_all_gather_axis must be 0 or 1, got {sampling_all_gather_axis}") |
|
|
