Add shared-parameter Ouro Megatron provider

vime_plugins/ouro/__init__.py

@@ -0,0 +1 @@
+"""Ouro recurrent training plugins (TP/PP/CP=1)."""

vime_plugins/ouro/model.py

@@ -0,0 +1,138 @@
+"""Shared-parameter Ouro for Megatron's optimizer, DDP and checkpoint contracts.
+
+The optional vllm-rlt dependency supplies checkpoint-compatible physical modules.
+Training uses differentiable SDPA, not the serving KV writes or a HF model wrapper.
+"""
+
+from functools import partial
+
+import torch
+import torch.nn.functional as F
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.transformer.module import MegatronModule
+from megatron.core.transformer.transformer_config import TransformerConfig
+from torch.utils.checkpoint import checkpoint
+from vllm_rlt.models.ouro import OuroDecoderLayer, OuroForCausalLM
+
+
+def decoder_forward(
+    layer: OuroDecoderLayer,
+    hidden: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> torch.Tensor:
+    attention = layer.self_attn
+    value = layer.input_layernorm(hidden)
+    shape = (hidden.shape[0], -1, attention.config.head_dim)
+    q, k, v = (projection(value).view(shape) for projection in (attention.q_proj, attention.k_proj, attention.v_proj))
+
+    def rotate(tensor: torch.Tensor) -> torch.Tensor:
+        first, second = tensor.chunk(2, dim=-1)
+        return tensor * cos + torch.cat((-second, first), dim=-1) * sin
+
+    q, k = rotate(q), rotate(k)
+    output = (
+        F.scaled_dot_product_attention(
+            q.transpose(0, 1).unsqueeze(0),
+            k.transpose(0, 1).unsqueeze(0),
+            v.transpose(0, 1).unsqueeze(0),
+            is_causal=True,
+            enable_gqa=True,
+        )
+        .squeeze(0)
+        .transpose(0, 1)
+        .reshape(hidden.shape[0], -1)
+    )
+    hidden = hidden + layer.input_layernorm_2(attention.o_proj(output))
+    return hidden + layer.post_attention_layernorm_2(layer.mlp(layer.post_attention_layernorm(hidden)))
+
+
+class OuroMegatronModel(MegatronModule):
+    def __init__(self, config: TransformerConfig, native: OuroForCausalLM, *, recompute: bool = False):
+        super().__init__(config)
+        if (config.tensor_model_parallel_size, config.pipeline_model_parallel_size, config.context_parallel_size) != (
+            1,
+            1,
+            1,
+        ):
+            raise ValueError("Ouro currently requires TP=PP=CP=1")
+        if config.sequence_parallel or config.gradient_accumulation_fusion:
+            raise ValueError("Ouro SDPA provider requires sequence parallel and grad fusion disabled")
+        self.ouro_config = native.config
+        self.model = native.model
+        self.lm_head = native.lm_head
+        self.requires_grad_(True)
+        self.model.early_exit_gate.requires_grad_(False)
+        self.pre_process = self.post_process = True
+        self.share_embeddings_and_output_weights = False
+        self.vp_stage = None
+        self.loop_budget = self.ouro_config.total_ut_steps
+        self.recompute = recompute
+        self.block_tokens = 0
+        self.recompute_block_tokens = 0
+
+    def set_input_tensor(self, input_tensor: list[torch.Tensor | None]) -> None:
+        if any(value is not None for value in input_tensor):
+            raise ValueError("Ouro does not accept pipeline input")
+
+    def set_loop_budget(self, loops: int) -> None:
+        if type(loops) is not int or not 1 <= loops <= self.ouro_config.total_ut_steps:
+            raise ValueError("Loop budget exceeds the checkpoint's supported depth")
+        self.loop_budget = loops
+
+    def _layer(
+        self, layer: OuroDecoderLayer, hidden: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+    ) -> torch.Tensor:
+        self.block_tokens += hidden.shape[0]
+        return decoder_forward(layer, hidden, cos, sin)
+
+    def _sequence(self, tokens: torch.Tensor) -> torch.Tensor:
+        hidden = self.model.embed_tokens(tokens)
+        positions = torch.arange(tokens.numel(), device=tokens.device)
+        cos, sin = self.model.rotary_emb(hidden, positions)
+        # Capture K in the forward, so later schedule changes cannot alter recomputation.
+        for _ in range(self.loop_budget):
+            for layer in self.model.layers:
+                if self.recompute and torch.is_grad_enabled():
+                    hidden = checkpoint(partial(self._layer, layer), hidden, cos, sin, use_reentrant=False)
+                else:
+                    hidden = self._layer(layer, hidden, cos, sin)
+            hidden = self.model.norm(hidden)
+        return self.lm_head(hidden)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        position_ids=None,
+        attention_mask=None,
+        labels=None,
+        packed_seq_params: PackedSeqParams | None = None,
+        loss_mask=None,
+    ) -> torch.Tensor:
+        if labels is not None or attention_mask is not None or position_ids is not None:
+            raise ValueError("Use VIME's masked RL loss and unmodified per-sequence positions")
+        if input_ids.ndim != 2 or input_ids.shape[0] != 1:
+            raise ValueError("Ouro expects one packed token stream")
+        if packed_seq_params is None:
+            return self._sequence(input_ids[0]).unsqueeze(0)
+        if packed_seq_params.qkv_format != "thd":
+            raise ValueError("Ouro expects thd packed sequences")
+        boundaries = packed_seq_params.cu_seqlens_q.tolist()
+        outputs = [
+            self._sequence(input_ids[0, start:end]) for start, end in zip(boundaries[:-1], boundaries[1:], strict=True)
+        ]
+        return torch.cat(outputs).unsqueeze(0)
+
+
+def model_provider(
+    pre_process: bool = True, post_process: bool = True, vp_stage: int | None = None
+) -> OuroMegatronModel:
+    from megatron.training import get_args
+    from megatron.training.arguments import core_transformer_config_from_args
+
+    if not pre_process or not post_process or vp_stage is not None:
+        raise ValueError("Ouro supports one non-pipelined model chunk")
+    args = get_args()
+    config = core_transformer_config_from_args(args)
+    native = OuroForCausalLM.from_pretrained(args.hf_checkpoint, dtype=config.params_dtype)
+    return OuroMegatronModel(config, native, recompute=args.recompute_granularity is not None)