| import math |
| from typing import List, Tuple |
| from dataclasses import dataclass |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| |
| |
| |
|
|
| def _trunc_normal_init(tensor: torch.Tensor, std: float = 1.0, |
| lower: float = -2.0, upper: float = 2.0) -> torch.Tensor: |
| with torch.no_grad(): |
| if std == 0: |
| tensor.zero_() |
| else: |
| sqrt2 = math.sqrt(2) |
| a = math.erf(lower / sqrt2) |
| b = math.erf(upper / sqrt2) |
| z = (b - a) / 2 |
| c = (2 * math.pi) ** -0.5 |
| pdf_u = c * math.exp(-0.5 * lower ** 2) |
| pdf_l = c * math.exp(-0.5 * upper ** 2) |
| comp_std = std / math.sqrt( |
| 1 - (upper * pdf_u - lower * pdf_l) / z |
| - ((pdf_u - pdf_l) / z) ** 2 |
| ) |
| tensor.uniform_(a, b) |
| tensor.erfinv_() |
| tensor.mul_(sqrt2 * comp_std) |
| tensor.clip_(lower * comp_std, upper * comp_std) |
| return tensor |
|
|
|
|
| def _find_multiple(a: int, b: int) -> int: |
| return (-(a // -b)) * b |
|
|
|
|
| def rms_norm(hidden_states: torch.Tensor, variance_epsilon: float) -> torch.Tensor: |
| input_dtype = hidden_states.dtype |
| hidden_states = hidden_states.to(torch.float32) |
| variance = hidden_states.square().mean(-1, keepdim=True) |
| hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon) |
| return hidden_states.to(input_dtype) |
|
|
|
|
| class CastedLinear(nn.Module): |
| def __init__(self, in_features: int, out_features: int, bias: bool): |
| super().__init__() |
| self.weight = nn.Parameter( |
| _trunc_normal_init(torch.empty(out_features, in_features), |
| std=1.0 / (in_features ** 0.5)) |
| ) |
| self.bias = nn.Parameter(torch.zeros(out_features)) if bias else None |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return F.linear( |
| x, self.weight.to(x.dtype), |
| bias=self.bias.to(x.dtype) if self.bias is not None else None, |
| ) |
|
|
|
|
| class SwiGLU(nn.Module): |
| def __init__(self, hidden_size: int, expansion: float): |
| super().__init__() |
| inter = _find_multiple(round(expansion * hidden_size * 2 / 3), 256) |
| self.gate_up_proj = CastedLinear(hidden_size, inter * 2, bias=False) |
| self.down_proj = CastedLinear(inter, hidden_size, bias=False) |
|
|
| def forward(self, x): |
| gate, up = self.gate_up_proj(x).chunk(2, dim=-1) |
| return self.down_proj(F.silu(gate) * up) |
|
|
|
|
| |
| |
| |
|
|
| class RTMBlock(nn.Module): |
|
|
| def __init__(self, hidden_size: int, seq_len: int, expansion: float, |
| rms_norm_eps: float = 1e-5): |
| super().__init__() |
| self.mlp_t = SwiGLU(hidden_size=seq_len, expansion=expansion) |
| self.mlp = SwiGLU(hidden_size=hidden_size, expansion=expansion) |
| self.norm_eps = rms_norm_eps |
|
|
| def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| hidden_states = hidden_states.transpose(1, 2) |
| hidden_states = rms_norm( |
| hidden_states + self.mlp_t(hidden_states), |
| variance_epsilon=self.norm_eps, |
| ) |
| hidden_states = hidden_states.transpose(1, 2) |
| hidden_states = rms_norm( |
| hidden_states + self.mlp(hidden_states), |
| variance_epsilon=self.norm_eps, |
| ) |
| return hidden_states |
|
|
|
|
| class RTMReasoningModule(nn.Module): |
|
|
| def __init__(self, layers: List[RTMBlock]): |
| super().__init__() |
| self.layers = nn.ModuleList(layers) |
|
|
| def forward(self, hidden_states: torch.Tensor, |
| input_injection: torch.Tensor) -> torch.Tensor: |
| hidden_states = hidden_states + input_injection |
| for layer in self.layers: |
| hidden_states = layer(hidden_states) |
| return hidden_states |
|
|
|
|
| def _make_level(num_layers: int, hidden_size: int, seq_len: int, |
| expansion: float, rms_norm_eps: float) -> RTMReasoningModule: |
| return RTMReasoningModule([ |
| RTMBlock(hidden_size, seq_len, expansion, rms_norm_eps) |
| for _ in range(num_layers) |
| ]) |
|
|
|
|
| |
| |
| |
|
|
| @dataclass |
| class RTMCarry: |
| z_H: torch.Tensor |
| z_L: torch.Tensor |
|
|
|
|
| class RTMInner(nn.Module): |
|
|
| def __init__(self, hidden_size: int, expansion: float, |
| H_cycles: int, L_cycles: int, H_layers: int, L_layers: int, |
| num_tokens: int, |
| with_grad: bool = False, |
| cycle_noise_std: float = 0.0, |
| rms_norm_eps: float = 1e-5, |
| forward_dtype: str = "float32"): |
| super().__init__() |
|
|
| self.hidden_size = hidden_size |
| self.H_cycles = H_cycles |
| self.L_cycles = L_cycles |
| self.with_grad = with_grad |
| self.cycle_noise_std = max(0.0, float(cycle_noise_std)) |
| self.forward_dtype = getattr(torch, forward_dtype) |
|
|
| self.total_seq_len = max(1, int(num_tokens)) |
|
|
| self.L_level = _make_level(L_layers, hidden_size, self.total_seq_len, |
| expansion, rms_norm_eps) |
| self.H_init = nn.Buffer( |
| _trunc_normal_init(torch.empty(hidden_size, dtype=self.forward_dtype), std=1), |
| persistent=True, |
| ) |
| self.L_init = nn.Parameter( |
| _trunc_normal_init(torch.empty(hidden_size, dtype=self.forward_dtype), std=1) |
| ) |
|
|
| def empty_carry(self, batch_size: int, device=None) -> RTMCarry: |
| if device is None: |
| device = self.H_init.device |
| return RTMCarry( |
| z_H=self.H_init.unsqueeze(0).unsqueeze(0).expand( |
| batch_size, self.total_seq_len, -1), |
| z_L=self.L_init.unsqueeze(0).unsqueeze(0).expand( |
| batch_size, self.total_seq_len, -1), |
| ) |
|
|
| def forward(self, carry: RTMCarry, z_H_init: torch.Tensor |
| ) -> Tuple[RTMCarry, torch.Tensor, List[torch.Tensor]]: |
| z_H, z_L = carry.z_H, carry.z_L |
| intermediates: List[torch.Tensor] = [] |
|
|
| if self.with_grad: |
| for _ in range(self.H_cycles): |
| for _ in range(self.L_cycles): |
| z_L = self.L_level(z_L, z_H + z_H_init) |
| z_H = self.L_level(z_H, z_L) |
| if self.training and self.cycle_noise_std > 0: |
| z_H = z_H + torch.randn_like(z_H) * self.cycle_noise_std |
| intermediates.append(z_H) |
| else: |
| with torch.no_grad(): |
| for _ in range(self.H_cycles - 1): |
| for _ in range(self.L_cycles): |
| z_L = self.L_level(z_L, z_H + z_H_init) |
| z_H = self.L_level(z_H, z_L) |
| if self.training and self.cycle_noise_std > 0: |
| z_H = z_H + torch.randn_like(z_H) * self.cycle_noise_std |
| intermediates.append(z_H) |
|
|
| for _ in range(self.L_cycles): |
| z_L = self.L_level(z_L, z_H + z_H_init) |
| z_H = self.L_level(z_H, z_L) |
| if self.training and self.cycle_noise_std > 0: |
| z_H = z_H + torch.randn_like(z_H) * self.cycle_noise_std |
| intermediates.append(z_H) |
|
|
| new_carry = RTMCarry(z_H=z_H.detach(), z_L=z_L.detach()) |
| return new_carry, z_H, intermediates |
|
|
| class _PixelNorm(nn.Module): |
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return x * (x.square().mean(dim=1, keepdim=True) + 1e-8).rsqrt() |
|
|
|
|
| class _EqualLinear(nn.Module): |
| def __init__(self, in_dim: int, out_dim: int): |
| super().__init__() |
| self.linear = nn.Linear(in_dim, out_dim) |
| self.linear.bias.data.zero_() |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.linear(x) |
|
|
|
|
| def _coerce_noise_2d(input_tensor: torch.Tensor, code_dim: int, mapper_name: str) -> torch.Tensor: |
| """Accept [B, D] or [B, 1, D] and return [B, D].""" |
| if input_tensor.dim() == 2 and input_tensor.shape[1] == code_dim: |
| return input_tensor |
| if input_tensor.dim() == 3 and input_tensor.shape[1] == 1 and input_tensor.shape[2] == code_dim: |
| return input_tensor.squeeze(1) |
| raise ValueError( |
| f"{mapper_name} expected input shape [B, {code_dim}] or [B, 1, {code_dim}], " |
| f"got {tuple(input_tensor.shape)}" |
| ) |
|
|
|
|
| class RTMMappingNetwork(nn.Module): |
| def __init__(self, code_dim: int, |
| num_tokens: int = 1, |
| H_cycles: int = 1, L_cycles: int = 1, |
| H_layers: int = 2, L_layers: int = 2, |
| hidden_size: int = 256, |
| expansion: float = 4.0, refinement_steps: int = 1, |
| with_grad: bool = False, |
| cycle_noise_std: float = 0.0, |
| rms_norm_eps: float = 1e-5, |
| forward_dtype: str = "float32"): |
| super().__init__() |
| self.code_dim = code_dim |
| self.refinement_steps = max(1, refinement_steps) |
|
|
| total_seq_len = max(1, int(num_tokens)) |
|
|
| if hidden_size <= 0: |
| assert code_dim % total_seq_len == 0, ( |
| f"code_dim={code_dim} must be divisible by num_tokens={total_seq_len} " |
| f"when hidden_size is auto" |
| ) |
| hidden_size = code_dim // total_seq_len |
|
|
| self.hidden_size = hidden_size |
|
|
| self.trm = RTMInner( |
| hidden_size=hidden_size, |
| expansion=expansion, |
| H_cycles=H_cycles, |
| L_cycles=L_cycles, |
| H_layers=H_layers, |
| L_layers=L_layers, |
| num_tokens=num_tokens, |
| with_grad=with_grad, |
| cycle_noise_std=cycle_noise_std, |
| rms_norm_eps=rms_norm_eps, |
| forward_dtype=forward_dtype, |
| ) |
|
|
| seq_len = self.trm.total_seq_len |
| capacity = seq_len * hidden_size |
|
|
| self.pixel_norm = _PixelNorm() |
| self.mapper_direct = (capacity == code_dim) |
| if not self.mapper_direct: |
| self.z_to_seq = _EqualLinear(code_dim, capacity) |
| self.seq_to_w = _EqualLinear(capacity, code_dim) |
|
|
| def _z_H_to_w(self, z_H: torch.Tensor) -> torch.Tensor: |
| flat = z_H.flatten(start_dim=1) |
| if self.mapper_direct: |
| return flat |
| return self.seq_to_w(flat) |
|
|
| def forward(self, input): |
| if isinstance(input, (list, tuple)): |
| input = input[0] |
|
|
| input = _coerce_noise_2d(input, self.code_dim, self.__class__.__name__) |
| B = input.shape[0] |
| seq_len = self.trm.total_seq_len |
|
|
| z_norm = self.pixel_norm(input) |
| if self.mapper_direct: |
| z_seq = z_norm.view(B, seq_len, self.hidden_size) |
| else: |
| z_seq = self.z_to_seq(z_norm).view(B, seq_len, self.hidden_size) |
|
|
| carry = self.trm.empty_carry(B, device=input.device) |
| for _ in range(self.refinement_steps): |
| carry, z_H_out, _ = self.trm(carry, z_H_init=z_seq) |
|
|
| return [self._z_H_to_w(z_H_out)] |
|
|
| def forward_w_trajectory(self, input): |
| """Decode every intermediate H state into a style vector. |
| |
| Returns ``[w_proj, w_after_cycle_1, ...]`` so callers can visualize |
| how w is refined across the H/L cycles. |
| """ |
| if isinstance(input, (list, tuple)): |
| input = input[0] |
|
|
| input = _coerce_noise_2d(input, self.code_dim, self.__class__.__name__) |
| B = input.shape[0] |
| seq_len = self.trm.total_seq_len |
|
|
| z_norm = self.pixel_norm(input) |
| if self.mapper_direct: |
| z_seq = z_norm.view(B, seq_len, self.hidden_size) |
| else: |
| z_seq = self.z_to_seq(z_norm).view(B, seq_len, self.hidden_size) |
|
|
| trajectory = [self._z_H_to_w(z_seq)] |
| carry = self.trm.empty_carry(B, device=input.device) |
| for _ in range(self.refinement_steps): |
| carry, _z_H_out, intermediates = self.trm(carry, z_H_init=z_seq) |
| for zh in intermediates: |
| trajectory.append(self._z_H_to_w(zh)) |
| return trajectory |
|
|
