| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """The ``WeightSource`` for ``rlhf_sharded_rdt_small_ep.py``. |
| 4 | |
| 5 | A trainer's producer job is to present its weights under the names the inference |
| 6 | side expects, which for vLLM means HF CHECKPOINT names. Transformers does not |
| 7 | store MoE experts that way -- it fuses each layer's into ``[E, ...]`` tensors -- |
| 8 | so that is the one conversion this file does. Kept beside the example rather than |
| 9 | inside it so the example stays about the sync itself, alongside |
| 10 | ``rdt_vllm_serve.py``. |
| 11 | """ |
| 12 | |
| 13 | import regex as re |
| 14 | import torch |
| 15 | |
| 16 | from vllm.distributed.weight_transfer import ParamMeta, WeightSource |
| 17 | from vllm.distributed.weight_transfer.base import materialize_full_tensor |
| 18 | |
| 19 | # Transformers fuses a layer's experts into ``[E, ...]`` tensors; the checkpoint |
| 20 | # stores them per expert. This matches the fused parameter names. |
| 21 | _FUSED_EXPERT_RE = re.compile(r"^(.*\.experts)\.(gate_up_proj|down_proj)$") |
| 22 | |
| 23 | |
| 24 | class CheckpointNameSource(WeightSource): |
| 25 | """`WeightSource` over an FSDP2 model that publishes HF CHECKPOINT names. |
| 26 | |
| 27 | Transformers keeps each layer's experts fused as ``[E, ...]`` tensors, but |
| 28 | vLLM's MoE loaders are written against the per-expert checkpoint entries -- |
| 29 | only ``qwen{2,3}_moe`` also accept the fused names, so publishing them |
| 30 | straight from ``named_parameters()`` limits the example to one family and |
| 31 | fails the consumer's bake on ``experts.gate_up_proj`` everywhere else. |
| 32 | Splitting them back is also what a real trainer does, since converting an |
| 33 | internal layout to checkpoint names is the normal producer job. |
| 34 | |
| 35 | Every rank publishes the whole model: fused params are all-gathered and then |
| 36 | sliced per expert. Ownership stays uniform (no ``held_names``) -- serving |
| 37 | only the experts a rank holds is a real trainer's concern, not an example's. |
| 38 | """ |
| 39 | |
| 40 | def __init__(self, module: torch.nn.Module) -> None: |
| 41 | self._module = module |
| 42 | |
| 43 | @staticmethod |
| 44 | def _expand(name: str, param: torch.Tensor) -> list[tuple[str, tuple]]: |
| 45 | """The checkpoint entries ``name`` contributes, as (name, shape). |
| 46 | |
| 47 | ``gate_up_proj`` is ``[E, 2I, H]`` with gate in rows ``:I`` and up in |
| 48 | ``I:``; ``down_proj`` is ``[E, H, I]``. Anything else is already a |
| 49 | checkpoint name and passes through. |
| 50 | """ |
| 51 | m = _FUSED_EXPERT_RE.match(name) |
| 52 | if m is None: |
| 53 | return [(name, tuple(param.shape))] |
| 54 | prefix, kind = m.group(1), m.group(2) |
| 55 | experts, rows, cols = param.shape |
| 56 | if kind == "down_proj": |
| 57 | return [ |
| 58 | (f"{prefix}.{e}.down_proj.weight", (rows, cols)) for e in range(experts) |
| 59 | ] |
| 60 | half = rows // 2 |
| 61 | return [ |
| 62 | (f"{prefix}.{e}.{proj}_proj.weight", (half, cols)) |
| 63 | for e in range(experts) |
| 64 | for proj in ("gate", "up") |
| 65 | ] |
| 66 | |
| 67 | def metadata(self) -> list[ParamMeta]: |
| 68 | return [ |
| 69 | ParamMeta(entry_name, param.dtype, shape) |
| 70 | for name, param in self._module.named_parameters() |
| 71 | for entry_name, shape in self._expand(name, param) |
| 72 | ] |
| 73 | |
| 74 | def __iter__(self): |
| 75 | for name, param in self._module.named_parameters(): |
| 76 | m = _FUSED_EXPERT_RE.match(name) |
| 77 | if m is None: |
| 78 | yield name, materialize_full_tensor(param) |
| 79 | continue |
| 80 | # One gather per FUSED param, not per expert: the views below are |
| 81 | # into that one tensor, and every expert of a layer rides in the |
| 82 | # same gather group, so it stays resident exactly as long as needed. |
| 83 | prefix, kind = m.group(1), m.group(2) |
| 84 | full = materialize_full_tensor(param) |
| 85 | if kind == "down_proj": |
| 86 | for e in range(full.shape[0]): |
| 87 | yield f"{prefix}.{e}.down_proj.weight", full[e] |
| 88 | else: |
| 89 | half = full.shape[1] // 2 |
| 90 | for e in range(full.shape[0]): |
| 91 | yield f"{prefix}.{e}.gate_proj.weight", full[e, :half, :] |
| 92 | yield f"{prefix}.{e}.up_proj.weight", full[e, half:, :] |