| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch.utils.checkpoint import checkpoint |
|
|
| from .attention import GQAAttentionCache, GatedGroupedQueryAttention |
| from .attn_res import ( |
| SlicedDepthAttentionRouter, |
| SublayerDepthRouters, |
| sliced_block_accumulate, |
| sliced_phase1_from_logits, |
| sliced_phase2_merge, |
| sliced_source_logits, |
| ) |
| from .config import YunmoNextConfig |
| from .kda_runtime import configure_kda_mixer_runtime |
| from .packed import document_position_ids |
| from .training_engine import TrainingEngineConfig |
|
|
|
|
| ATTENTION_RESIDUAL_QUERY_INIT_SEED = 17 |
| ATTENTION_RESIDUAL_QUERY_INIT_OFFSET = 7_900_031 |
| ATTENTION_RESIDUAL_QUERY_INIT_STD = 0.02 |
| ATTENTION_RESIDUAL_ORTHOGONAL_QUERY_DEPTH = 24 |
| ACTIVATION_CHECKPOINT_POLICY_FULL = "full" |
| ACTIVATION_CHECKPOINT_POLICIES = frozenset({ACTIVATION_CHECKPOINT_POLICY_FULL}) |
|
|
|
|
| def orthogonal_attention_residual_query_bank( |
| query_count: int, |
| key_dim: int, |
| *, |
| device: torch.device | str = "cpu", |
| dtype: torch.dtype = torch.float32, |
| ) -> torch.Tensor: |
| """Return the adopted private-seed orthogonal 24-layer query bank. |
| |
| The first row is generated for lineage stability but remains zero in the |
| model: the first pre-mixer read has one completed source, so that query is |
| structurally unidentifiable. A private CPU generator prevents this |
| initialization from advancing the common model RNG stream. |
| """ |
|
|
| if query_count <= 0 or key_dim <= 0 or query_count > key_dim: |
| raise ValueError( |
| "orthogonal Attention Residual query bank requires " |
| "0 < queries <= key_dim" |
| ) |
| generator = torch.Generator(device="cpu") |
| generator.manual_seed( |
| ( |
| ATTENTION_RESIDUAL_QUERY_INIT_SEED |
| + ATTENTION_RESIDUAL_QUERY_INIT_OFFSET |
| ) |
| % (2**63 - 1) |
| ) |
| basis = torch.randn( |
| key_dim, |
| query_count, |
| generator=generator, |
| dtype=torch.float32, |
| ) |
| rows = torch.linalg.qr(basis, mode="reduced").Q.T.contiguous() |
| rows.mul_((key_dim**0.5) * ATTENTION_RESIDUAL_QUERY_INIT_STD) |
| return rows.to(device=device, dtype=dtype) |
|
|
|
|
| def _build_recurrent_mixer( |
| config: YunmoNextConfig, |
| engine: TrainingEngineConfig, |
| layer_idx: int, |
| ) -> nn.Module: |
| try: |
| from fla.layers import KimiDeltaAttention |
| except ImportError as exc: |
| if not torch.cuda.is_available(): |
| if engine.recurrent_segment_tokens: |
| raise RuntimeError( |
| "serial recurrent execution requires the CUDA KDA runtime" |
| ) from exc |
| from .reference_recurrent import ReferenceKimiDeltaAttention |
|
|
| return ReferenceKimiDeltaAttention( |
| hidden_size=config.hidden_size, |
| head_dim=config.recurrent_head_dim, |
| num_heads=config.recurrent_num_heads, |
| expand_v=config.recurrent_expand_v, |
| use_short_conv=True, |
| conv_size=config.short_conv_size, |
| layer_idx=layer_idx, |
| norm_eps=config.rms_norm_eps, |
| ) |
| raise ImportError( |
| "Yunmo-Next recurrent layers require the FLA CUDA environment" |
| ) from exc |
|
|
| mixer = KimiDeltaAttention( |
| hidden_size=config.hidden_size, |
| head_dim=config.recurrent_head_dim, |
| num_heads=config.recurrent_num_heads, |
| expand_v=config.recurrent_expand_v, |
| use_short_conv=True, |
| conv_size=config.short_conv_size, |
| mode="chunk", |
| layer_idx=layer_idx, |
| norm_eps=config.rms_norm_eps, |
| ) |
| return configure_kda_mixer_runtime(mixer, engine) |
|
|
|
|
| class _SerialChunkedSwiGLU(torch.autograd.Function): |
| @staticmethod |
| def forward( |
| ctx, |
| hidden_states: torch.Tensor, |
| gate_weight: torch.Tensor, |
| up_weight: torch.Tensor, |
| down_weight: torch.Tensor, |
| chunk_size: int, |
| ) -> torch.Tensor: |
| ctx.chunk_size = chunk_size |
| ctx.save_for_backward(hidden_states, gate_weight, up_weight, down_weight) |
| return torch.cat( |
| [ |
| F.linear( |
| F.silu(F.linear(chunk, gate_weight)) |
| * F.linear(chunk, up_weight), |
| down_weight, |
| ) |
| for chunk in hidden_states.split(chunk_size, dim=-2) |
| ], |
| dim=-2, |
| ) |
|
|
| @staticmethod |
| def backward(ctx, grad_output: torch.Tensor): |
| hidden_states, gate_weight, up_weight, down_weight = ctx.saved_tensors |
| grad_hidden_chunks: list[torch.Tensor] = [] |
| accumulator_dtype = ( |
| torch.float64 if gate_weight.dtype == torch.float64 else torch.float32 |
| ) |
| grad_gate = torch.zeros_like(gate_weight, dtype=accumulator_dtype) |
| grad_up = torch.zeros_like(up_weight, dtype=accumulator_dtype) |
| grad_down = torch.zeros_like(down_weight, dtype=accumulator_dtype) |
| hidden_chunks = hidden_states.split(ctx.chunk_size, dim=-2) |
| output_grad_chunks = grad_output.split(ctx.chunk_size, dim=-2) |
| create_graph = torch.is_grad_enabled() |
| for hidden_chunk, output_grad_chunk in zip( |
| hidden_chunks, output_grad_chunks, strict=True |
| ): |
| with torch.enable_grad(): |
| chunk = hidden_chunk.detach().requires_grad_(True) |
| gate = gate_weight.detach().requires_grad_(True) |
| up = up_weight.detach().requires_grad_(True) |
| down = down_weight.detach().requires_grad_(True) |
| output = F.linear( |
| F.silu(F.linear(chunk, gate)) * F.linear(chunk, up), |
| down, |
| ) |
| chunk_grad, gate_grad, up_grad, down_grad = torch.autograd.grad( |
| output, |
| (chunk, gate, up, down), |
| output_grad_chunk, |
| create_graph=create_graph, |
| ) |
| grad_hidden_chunks.append(chunk_grad) |
| grad_gate.add_(gate_grad.to(accumulator_dtype)) |
| grad_up.add_(up_grad.to(accumulator_dtype)) |
| grad_down.add_(down_grad.to(accumulator_dtype)) |
| return ( |
| torch.cat(grad_hidden_chunks, dim=-2), |
| grad_gate.to(gate_weight.dtype), |
| grad_up.to(up_weight.dtype), |
| grad_down.to(down_weight.dtype), |
| None, |
| ) |
|
|
|
|
| class SwiGLU(nn.Module): |
| def __init__( |
| self, |
| hidden_size: int, |
| intermediate_size: int, |
| engine: TrainingEngineConfig, |
| ) -> None: |
| super().__init__() |
| self.sequence_chunk_size = engine.feed_forward_chunk_tokens |
| self.sequence_serial_backward = engine.feed_forward_serial_backward |
| self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) |
| self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) |
| self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) |
|
|
| def _project(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| gate = self.gate_proj(hidden_states) |
| up = self.up_proj(hidden_states) |
| activated = F.silu(gate) * up |
| return self.down_proj(activated) |
|
|
| def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| chunk_size = self.sequence_chunk_size |
| if chunk_size <= 0 or hidden_states.shape[-2] <= chunk_size: |
| return self._project(hidden_states) |
| if self.sequence_serial_backward: |
| return _SerialChunkedSwiGLU.apply( |
| hidden_states, |
| self.gate_proj.weight, |
| self.up_proj.weight, |
| self.down_proj.weight, |
| chunk_size, |
| ) |
| return torch.cat( |
| [ |
| self._project(chunk) |
| for chunk in hidden_states.split(chunk_size, dim=-2) |
| ], |
| dim=-2, |
| ) |
|
|
|
|
| class YunmoBlock(nn.Module): |
| def __init__( |
| self, |
| config: YunmoNextConfig, |
| engine: TrainingEngineConfig, |
| layer_idx: int, |
| ) -> None: |
| super().__init__() |
| self.layer_idx = layer_idx |
| self.input_norm = nn.RMSNorm( |
| config.hidden_size, |
| eps=config.rms_norm_eps, |
| ) |
| self.post_mixer_norm = nn.RMSNorm( |
| config.hidden_size, |
| eps=config.rms_norm_eps, |
| ) |
| self.mixer = ( |
| GatedGroupedQueryAttention(config, layer_idx) |
| if config.is_attention_layer(layer_idx) |
| else _build_recurrent_mixer(config, engine, layer_idx) |
| ) |
| self.mlp = SwiGLU( |
| config.hidden_size, |
| config.intermediate_size, |
| engine, |
| ) |
| self.depth_routers = SublayerDepthRouters( |
| config.hidden_size, |
| config.rms_norm_eps, |
| ) |
|
|
| def mixer_output( |
| self, |
| hidden_states: torch.Tensor, |
| position_ids: torch.Tensor, |
| cu_seqlens: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| """Return one mixer transformation before its residual write.""" |
|
|
| normalized = self.input_norm(hidden_states) |
| if isinstance(self.mixer, GatedGroupedQueryAttention): |
| return self.mixer(normalized, position_ids, cu_seqlens) |
| return self.mixer(normalized, cu_seqlens=cu_seqlens)[0] |
|
|
| def mlp_output(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| """Return one feed-forward transformation before its residual write.""" |
|
|
| return self.mlp(self.post_mixer_norm(hidden_states)) |
|
|
| def mixer_output_cached( |
| self, |
| hidden_states: torch.Tensor, |
| position_ids: torch.Tensor, |
| *, |
| fla_cache: object, |
| attention_cache: GQAAttentionCache | None, |
| ) -> tuple[torch.Tensor, GQAAttentionCache | None]: |
| normalized = self.input_norm(hidden_states) |
| if isinstance(self.mixer, GatedGroupedQueryAttention): |
| if attention_cache is not None and not isinstance( |
| attention_cache, |
| GQAAttentionCache, |
| ): |
| raise TypeError("wrong cache type for grouped-query attention") |
| return self.mixer.forward_cached( |
| normalized, |
| position_ids, |
| attention_cache, |
| ) |
| mixed, _, _ = self.mixer( |
| normalized, |
| past_key_values=fla_cache, |
| use_cache=True, |
| ) |
| return mixed, attention_cache |
|
|
|
|
| @dataclass |
| class YunmoNextOutput: |
| loss: torch.Tensor | None = None |
| logits: torch.Tensor | None = None |
| hidden_states: torch.Tensor | None = None |
|
|
|
|
| @dataclass |
| class YunmoNextCUDAState: |
| """Bounded KDA/local-attention state plus sparse global KV.""" |
|
|
| fla_cache: object |
| attention: dict[int, GQAAttentionCache] |
| seen_tokens: int = 0 |
|
|
|
|
| class YunmoNextForCausalLM(nn.Module): |
| def __init__( |
| self, |
| config: YunmoNextConfig, |
| training_engine: TrainingEngineConfig | None = None, |
| ) -> None: |
| super().__init__() |
| self.config = config |
| self.training_engine = ( |
| TrainingEngineConfig() |
| if training_engine is None |
| else training_engine |
| ) |
| if not isinstance(self.training_engine, TrainingEngineConfig): |
| raise TypeError("training_engine must be TrainingEngineConfig") |
| self.activation_checkpointing_runtime_policy = ( |
| ACTIVATION_CHECKPOINT_POLICY_FULL |
| ) |
| self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) |
| self.layers = nn.ModuleList( |
| [ |
| YunmoBlock(config, self.training_engine, layer_idx) |
| for layer_idx in range(config.num_hidden_layers) |
| ] |
| ) |
| self._enable_cell_echo_biases() |
| self.final_depth_router = SlicedDepthAttentionRouter( |
| config.hidden_size, |
| config.rms_norm_eps, |
| ) |
| self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| self.reset_parameters() |
| self._install_cell_phase_mlp_geometry() |
|
|
| def set_activation_checkpointing_runtime_policy( |
| self, |
| policy: str, |
| ) -> None: |
| if policy not in ACTIVATION_CHECKPOINT_POLICIES: |
| raise ValueError( |
| f"unsupported activation-checkpointing runtime policy: {policy}" |
| ) |
| self.activation_checkpointing_runtime_policy = policy |
|
|
| def _activation_checkpoint( |
| self, |
| function, |
| *args: torch.Tensor, |
| ): |
| return checkpoint(function, *args, use_reentrant=False) |
|
|
| @torch.no_grad() |
| def reset_parameters(self) -> None: |
| nn.init.normal_( |
| self.embed_tokens.weight, |
| mean=0.0, |
| std=self.config.initializer_range, |
| ) |
| for module in self.modules(): |
| if isinstance(module, nn.Linear): |
| nn.init.normal_( |
| module.weight, |
| mean=0.0, |
| std=self.config.initializer_range, |
| ) |
|
|
| residual_scale = self.config.residual_write_initialization_scale |
| for block in self.layers: |
| block.mlp.down_proj.weight.mul_(residual_scale) |
| output_projection = getattr(block.mixer, "o_proj", None) |
| if output_projection is not None: |
| output_projection.weight.mul_(residual_scale) |
|
|
| self._reset_attention_residual_queries() |
|
|
| @torch.no_grad() |
| def _install_cell_phase_mlp_geometry(self) -> None: |
| """Install layer-local widths without perturbing non-MLP RNG mapping. |
| |
| Every layer first constructs and initializes the common baseline width. |
| Differently shaped MLPs are then installed from private layer-local |
| random streams without moving the rest of the model's seed mapping. |
| """ |
|
|
| seed = int(torch.initial_seed()) |
| modulus = 2**63 - 1 |
| for layer_idx, (layer, width) in enumerate( |
| zip( |
| self.layers, |
| self.config.mlp_intermediate_sizes, |
| strict=True, |
| ) |
| ): |
| mlp = SwiGLU( |
| self.config.hidden_size, |
| width, |
| self.training_engine, |
| ) |
| if mlp.gate_proj.weight.device.type == "meta": |
| layer.mlp = mlp |
| continue |
| generator = torch.Generator(device=mlp.gate_proj.weight.device) |
| generator.manual_seed( |
| (seed * 1_000_000 + 3 + layer_idx) % modulus |
| ) |
| for projection in ( |
| mlp.gate_proj, |
| mlp.up_proj, |
| mlp.down_proj, |
| ): |
| nn.init.normal_( |
| projection.weight, |
| mean=0.0, |
| std=self.config.initializer_range, |
| generator=generator, |
| ) |
| mlp.down_proj.weight.mul_( |
| self.config.residual_write_initialization_scale |
| ) |
| layer.mlp = mlp |
|
|
| @torch.no_grad() |
| def _reset_attention_residual_queries(self) -> None: |
| if ( |
| self.config.num_hidden_layers |
| != ATTENTION_RESIDUAL_ORTHOGONAL_QUERY_DEPTH |
| ): |
| return |
| routers = self._sliced_router_bank() |
| expected = 2 * ATTENTION_RESIDUAL_ORTHOGONAL_QUERY_DEPTH + 1 |
| if len(routers) != expected or any( |
| router.key_dim != 64 for router in routers |
| ): |
| raise RuntimeError( |
| "orthogonal Attention Residual initialization requires " |
| "49 rank-64 queries" |
| ) |
| bank = orthogonal_attention_residual_query_bank( |
| len(routers), |
| routers[0].key_dim, |
| device=routers[0].query.device, |
| dtype=routers[0].query.dtype, |
| ) |
| routers[0].query.zero_() |
| for router, row in zip(routers[1:], bank[1:], strict=True): |
| router.query.copy_(row) |
|
|
| def parameter_count(self) -> int: |
| return sum(parameter.numel() for parameter in self.parameters()) |
|
|
| def _apply_depth_router( |
| self, |
| router: SlicedDepthAttentionRouter, |
| sources: list[torch.Tensor], |
| routing_trace: list[torch.Tensor] | None, |
| ) -> torch.Tensor: |
| if ( |
| self.training_engine.activation_checkpointing |
| and self.training |
| and routing_trace is None |
| ): |
| def route_checkpointed(*inputs: torch.Tensor) -> torch.Tensor: |
| routed = router(list(inputs)) |
| assert isinstance(routed, torch.Tensor) |
| return routed |
|
|
| return self._activation_checkpoint(route_checkpointed, *sources) |
| if routing_trace is None: |
| output = router(sources) |
| assert isinstance(output, torch.Tensor) |
| return output |
| output, weights = router(sources, return_weights=True) |
| routing_trace.append(weights.detach()) |
| return output |
|
|
| def _mixer_output( |
| self, |
| layer: YunmoBlock, |
| hidden_states: torch.Tensor, |
| position_ids: torch.Tensor, |
| cu_seqlens: torch.Tensor | None, |
| ) -> torch.Tensor: |
| if self.training_engine.activation_checkpointing and self.training: |
| return self._activation_checkpoint( |
| layer.mixer_output, |
| hidden_states, |
| position_ids, |
| cu_seqlens, |
| ) |
| return layer.mixer_output(hidden_states, position_ids, cu_seqlens) |
|
|
| def _mlp_output( |
| self, |
| layer: YunmoBlock, |
| hidden_states: torch.Tensor, |
| ) -> torch.Tensor: |
| if self.training_engine.activation_checkpointing and self.training: |
| return self._activation_checkpoint(layer.mlp_output, hidden_states) |
| return layer.mlp_output(hidden_states) |
|
|
| def _sliced_router_bank(self) -> list[SlicedDepthAttentionRouter]: |
| routers: list[SlicedDepthAttentionRouter] = [] |
| for layer in self.layers: |
| routers.extend( |
| ( |
| layer.depth_routers.pre_mixer, |
| layer.depth_routers.pre_mlp, |
| ) |
| ) |
| routers.append(self.final_depth_router) |
| return routers |
|
|
| def _enable_cell_echo_biases(self) -> None: |
| """Address each attention-first cell's completed GQA source.""" |
|
|
| block_size = self.config.attn_res_sublayers_per_block |
| period = self.config.recurrent_period |
| for cell_start in range(0, len(self.layers), period): |
| target_layer = cell_start + period - 1 |
| router = self.layers[target_layer].depth_routers.pre_mixer |
| gqa_sublayer = 2 * cell_start |
| preferred_source = 1 + gqa_sublayer // block_size |
| target_read = 2 * target_layer |
| if preferred_source * block_size > target_read: |
| raise RuntimeError( |
| "cell GQA source is not completed before Cell-Echo read" |
| ) |
| router.enable_source_bias(preferred_source) |
|
|
| def _forward_hidden_checkpointed_blocks( |
| self, |
| embedding: torch.Tensor, |
| position_ids: torch.Tensor, |
| cu_seqlens: torch.Tensor | None, |
| routers: list[SlicedDepthAttentionRouter], |
| ) -> torch.Tensor: |
| """Checkpoint each complete two-phase Attention Residual block.""" |
|
|
| block_size = self.config.attn_res_sublayers_per_block |
| total_sublayers = 2 * len(self.layers) |
| if block_size % 2 != 0 or total_sublayers % block_size != 0: |
| raise ValueError( |
| "block-checkpointed Attention Residual requires complete " |
| "even sublayer blocks" |
| ) |
|
|
| layers_per_block = block_size // 2 |
| completed_values = [embedding] |
| for layer_first in range(0, len(self.layers), layers_per_block): |
| layer_end = layer_first + layers_per_block |
| sublayer_first = 2 * layer_first |
| block_layers = tuple(self.layers[layer_first:layer_end]) |
| block_routers = tuple( |
| routers[sublayer_first : sublayer_first + block_size] |
| ) |
| source_count = len(completed_values) |
|
|
| def run_block( |
| *inputs: torch.Tensor, |
| block_layers: tuple[YunmoBlock, ...] = block_layers, |
| block_routers: tuple[SlicedDepthAttentionRouter, ...] = block_routers, |
| source_count: int = source_count, |
| ) -> torch.Tensor: |
| checkpoint_sources = list(inputs[:source_count]) |
| logits = torch.stack( |
| [ |
| sliced_source_logits( |
| source, |
| block_routers, |
| source_index=source_index, |
| ) |
| for source_index, source in enumerate(checkpoint_sources) |
| ], |
| dim=1, |
| ) |
| outputs, lses, _ = sliced_phase1_from_logits( |
| checkpoint_sources, |
| logits, |
| routers=block_routers, |
| ) |
| phase_outputs = outputs.unbind(0) |
| phase_lses = lses.unbind(0) |
|
|
| partial_block: torch.Tensor | None = None |
| local_sublayer = 0 |
| for layer in block_layers: |
| mixer_input = phase_outputs[local_sublayer] |
| if partial_block is not None: |
| mixer_input, _ = sliced_phase2_merge( |
| block_routers[local_sublayer], |
| partial_block, |
| mixer_input, |
| phase_lses[local_sublayer], |
| ) |
| mixer_delta = layer.mixer_output( |
| mixer_input, |
| position_ids, |
| cu_seqlens, |
| ) |
| partial_block = sliced_block_accumulate( |
| partial_block, |
| mixer_delta, |
| ) |
| local_sublayer += 1 |
|
|
| mlp_input, _ = sliced_phase2_merge( |
| block_routers[local_sublayer], |
| partial_block, |
| phase_outputs[local_sublayer], |
| phase_lses[local_sublayer], |
| ) |
| partial_block = sliced_block_accumulate( |
| partial_block, |
| layer.mlp_output(mlp_input), |
| ) |
| local_sublayer += 1 |
|
|
| if partial_block is None or local_sublayer != len(block_routers): |
| raise RuntimeError( |
| "Attention Residual block did not close exactly" |
| ) |
| return partial_block |
|
|
| completed_values.append( |
| self._activation_checkpoint(run_block, *completed_values) |
| ) |
|
|
| hidden_states = self._apply_depth_router( |
| self.final_depth_router, |
| completed_values, |
| None, |
| ) |
| return self.norm(hidden_states) |
|
|
| def _forward_hidden_two_phase_attention_residual( |
| self, |
| embedding: torch.Tensor, |
| position_ids: torch.Tensor, |
| cu_seqlens: torch.Tensor | None, |
| layer_trace: list[torch.Tensor] | None, |
| routing_trace: list[torch.Tensor] | None, |
| ) -> torch.Tensor: |
| routers = self._sliced_router_bank() |
| total_sublayers = 2 * len(self.layers) |
| block_size = self.config.attn_res_sublayers_per_block |
| if ( |
| self.training_engine.activation_checkpointing |
| and self.training |
| and layer_trace is None |
| and routing_trace is None |
| and block_size % 2 == 0 |
| and total_sublayers % block_size == 0 |
| ): |
| return self._forward_hidden_checkpointed_blocks( |
| embedding, |
| position_ids, |
| cu_seqlens, |
| routers, |
| ) |
|
|
| completed_values = [embedding] |
| completed_logits = [ |
| sliced_source_logits(embedding, routers, source_index=0) |
| ] |
| completed_first_reads = [0] |
| partial_block: torch.Tensor | None = None |
| phase_first = -1 |
| phase_end = -1 |
| phase_outputs: tuple[torch.Tensor, ...] = () |
| phase_lses: tuple[torch.Tensor, ...] = () |
| phase_weights: tuple[torch.Tensor, ...] | None = None |
|
|
| def invalidate_phase() -> None: |
| nonlocal phase_first, phase_end, phase_outputs, phase_lses, phase_weights |
| phase_first = -1 |
| phase_end = -1 |
| phase_outputs = () |
| phase_lses = () |
| phase_weights = None |
|
|
| def append_completed(source: torch.Tensor, first_read: int) -> None: |
| source_index = len(completed_values) |
| completed_values.append(source) |
| completed_logits.append( |
| sliced_source_logits( |
| source, |
| routers[first_read:], |
| source_index=source_index, |
| ) |
| ) |
| completed_first_reads.append(first_read) |
| invalidate_phase() |
|
|
| def ensure_phase(read_idx: int) -> None: |
| nonlocal phase_first, phase_end, phase_outputs, phase_lses, phase_weights |
| if phase_first <= read_idx < phase_end: |
| return |
| phase_first = read_idx |
| phase_end = ( |
| min( |
| ((read_idx // block_size) + 1) * block_size, |
| total_sublayers, |
| ) |
| if read_idx < total_sublayers |
| else total_sublayers + 1 |
| ) |
| logits = torch.stack( |
| [ |
| source_logits[ |
| phase_first - first_read : phase_end - first_read |
| ] |
| for source_logits, first_read in zip( |
| completed_logits, |
| completed_first_reads, |
| strict=True, |
| ) |
| ], |
| dim=1, |
| ) |
| outputs, lses, weights = sliced_phase1_from_logits( |
| completed_values, |
| logits, |
| routers=routers[phase_first:phase_end], |
| return_weights=routing_trace is not None, |
| ) |
| phase_outputs = outputs.unbind(0) |
| phase_lses = lses.unbind(0) |
| phase_weights = weights.unbind(0) if weights is not None else None |
|
|
| def read_residual(read_idx: int) -> torch.Tensor: |
| ensure_phase(read_idx) |
| phase_idx = read_idx - phase_first |
| output = phase_outputs[phase_idx] |
| weights = ( |
| phase_weights[phase_idx] if phase_weights is not None else None |
| ) |
| if partial_block is not None: |
| output, partial_probability = sliced_phase2_merge( |
| routers[read_idx], |
| partial_block, |
| output, |
| phase_lses[phase_idx], |
| ) |
| if weights is not None: |
| completed_mass = 1.0 - partial_probability |
| weights = torch.cat( |
| [ |
| weights * completed_mass.unsqueeze(0), |
| partial_probability.unsqueeze(0), |
| ], |
| dim=0, |
| ) |
| if routing_trace is not None: |
| assert weights is not None |
| routing_trace.append(weights.detach()) |
| return output |
|
|
| sublayer_idx = 0 |
| for layer in self.layers: |
| mixer_input = read_residual(sublayer_idx) |
| mixer_delta = self._mixer_output( |
| layer, |
| mixer_input, |
| position_ids, |
| cu_seqlens, |
| ) |
| partial_block = sliced_block_accumulate( |
| partial_block, |
| mixer_delta, |
| ) |
| sublayer_idx += 1 |
| if sublayer_idx % block_size == 0: |
| append_completed(partial_block, sublayer_idx) |
| partial_block = None |
|
|
| mlp_input = read_residual(sublayer_idx) |
| mlp_delta = self._mlp_output(layer, mlp_input) |
| partial_block = sliced_block_accumulate(partial_block, mlp_delta) |
| sublayer_idx += 1 |
| if layer_trace is not None: |
| layer_trace.append(partial_block) |
| if sublayer_idx % block_size == 0: |
| append_completed(partial_block, sublayer_idx) |
| partial_block = None |
|
|
| if partial_block is not None: |
| append_completed(partial_block, total_sublayers) |
| hidden_states = read_residual(total_sublayers) |
| return self.norm(hidden_states) |
|
|
| def forward_hidden( |
| self, |
| input_ids: torch.Tensor, |
| position_ids: torch.Tensor | None = None, |
| cu_seqlens: torch.Tensor | None = None, |
| *, |
| layer_trace: list[torch.Tensor] | None = None, |
| routing_trace: list[torch.Tensor] | None = None, |
| ) -> torch.Tensor: |
| return self._forward_hidden_from_embeddings( |
| self.embed_tokens(input_ids), |
| position_ids, |
| cu_seqlens, |
| layer_trace=layer_trace, |
| routing_trace=routing_trace, |
| ) |
|
|
| def _forward_hidden_from_embeddings( |
| self, |
| embeddings: torch.Tensor, |
| position_ids: torch.Tensor | None = None, |
| cu_seqlens: torch.Tensor | None = None, |
| *, |
| layer_trace: list[torch.Tensor] | None = None, |
| routing_trace: list[torch.Tensor] | None = None, |
| ) -> torch.Tensor: |
| """Run the adopted text stack from already-constructed embeddings. |
| |
| This is the only non-token entry seam. It adds no module or parameter |
| and leaves creation or transformation of embeddings to the caller, |
| keeping the ordinary token path independent of those concerns. |
| """ |
|
|
| if embeddings.ndim != 3: |
| raise ValueError( |
| "embeddings must have shape [batch, sequence, hidden]" |
| ) |
| batch, sequence, hidden_size = embeddings.shape |
| if hidden_size != self.config.hidden_size: |
| raise ValueError("embedding hidden size disagrees with model config") |
| if position_ids is None: |
| position_ids = ( |
| document_position_ids(cu_seqlens, sequence) |
| if cu_seqlens is not None |
| else torch.arange( |
| sequence, |
| device=embeddings.device, |
| ).expand(batch, -1) |
| ) |
| elif position_ids.shape != (batch, sequence): |
| raise ValueError( |
| "position IDs must match embedding batch and sequence" |
| ) |
| return self._forward_hidden_two_phase_attention_residual( |
| embeddings, |
| position_ids, |
| cu_seqlens, |
| layer_trace, |
| routing_trace, |
| ) |
|
|
| def _forward_hidden_two_phase_attention_residual_cached( |
| self, |
| embedding: torch.Tensor, |
| position_ids: torch.Tensor, |
| cache: YunmoNextCUDAState, |
| layer_trace: list[torch.Tensor] | None, |
| routing_trace: list[torch.Tensor] | None, |
| ) -> torch.Tensor: |
| routers = self._sliced_router_bank() |
| total_sublayers = 2 * len(self.layers) |
| block_size = self.config.attn_res_sublayers_per_block |
| completed_values = [embedding] |
| completed_logits = [ |
| sliced_source_logits(embedding, routers, source_index=0) |
| ] |
| completed_first_reads = [0] |
| partial_block: torch.Tensor | None = None |
| phase_first = -1 |
| phase_end = -1 |
| phase_outputs: tuple[torch.Tensor, ...] = () |
| phase_lses: tuple[torch.Tensor, ...] = () |
| phase_weights: tuple[torch.Tensor, ...] | None = None |
|
|
| def invalidate_phase() -> None: |
| nonlocal phase_first, phase_end, phase_outputs, phase_lses, phase_weights |
| phase_first = -1 |
| phase_end = -1 |
| phase_outputs = () |
| phase_lses = () |
| phase_weights = None |
|
|
| def append_completed(source: torch.Tensor, first_read: int) -> None: |
| source_index = len(completed_values) |
| completed_values.append(source) |
| completed_logits.append( |
| sliced_source_logits( |
| source, |
| routers[first_read:], |
| source_index=source_index, |
| ) |
| ) |
| completed_first_reads.append(first_read) |
| invalidate_phase() |
|
|
| def ensure_phase(read_idx: int) -> None: |
| nonlocal phase_first, phase_end, phase_outputs, phase_lses, phase_weights |
| if phase_first <= read_idx < phase_end: |
| return |
| phase_first = read_idx |
| phase_end = ( |
| min( |
| ((read_idx // block_size) + 1) * block_size, |
| total_sublayers, |
| ) |
| if read_idx < total_sublayers |
| else total_sublayers + 1 |
| ) |
| logits = torch.stack( |
| [ |
| source_logits[ |
| phase_first - first_read : phase_end - first_read |
| ] |
| for source_logits, first_read in zip( |
| completed_logits, |
| completed_first_reads, |
| strict=True, |
| ) |
| ], |
| dim=1, |
| ) |
| outputs, lses, weights = sliced_phase1_from_logits( |
| completed_values, |
| logits, |
| routers=routers[phase_first:phase_end], |
| return_weights=routing_trace is not None, |
| ) |
| phase_outputs = outputs.unbind(0) |
| phase_lses = lses.unbind(0) |
| phase_weights = weights.unbind(0) if weights is not None else None |
|
|
| def read_residual(read_idx: int) -> torch.Tensor: |
| ensure_phase(read_idx) |
| phase_idx = read_idx - phase_first |
| output = phase_outputs[phase_idx] |
| weights = ( |
| phase_weights[phase_idx] if phase_weights is not None else None |
| ) |
| if partial_block is not None: |
| output, partial_probability = sliced_phase2_merge( |
| routers[read_idx], |
| partial_block, |
| output, |
| phase_lses[phase_idx], |
| ) |
| if weights is not None: |
| weights = torch.cat( |
| [ |
| weights |
| * (1.0 - partial_probability).unsqueeze(0), |
| partial_probability.unsqueeze(0), |
| ], |
| dim=0, |
| ) |
| if routing_trace is not None: |
| assert weights is not None |
| routing_trace.append(weights.detach()) |
| return output |
|
|
| sublayer_idx = 0 |
| for layer_idx, layer in enumerate(self.layers): |
| mixer_input = read_residual(sublayer_idx) |
| mixer_delta, layer_attention_cache = layer.mixer_output_cached( |
| mixer_input, |
| position_ids, |
| fla_cache=cache.fla_cache, |
| attention_cache=cache.attention.get(layer_idx), |
| ) |
| if layer_attention_cache is not None: |
| cache.attention[layer_idx] = layer_attention_cache |
| partial_block = sliced_block_accumulate( |
| partial_block, |
| mixer_delta, |
| ) |
| sublayer_idx += 1 |
| if sublayer_idx % block_size == 0: |
| append_completed(partial_block, sublayer_idx) |
| partial_block = None |
|
|
| mlp_input = read_residual(sublayer_idx) |
| mlp_delta = layer.mlp_output(mlp_input) |
| partial_block = sliced_block_accumulate(partial_block, mlp_delta) |
| sublayer_idx += 1 |
| if layer_trace is not None: |
| layer_trace.append(partial_block) |
| if sublayer_idx % block_size == 0: |
| append_completed(partial_block, sublayer_idx) |
| partial_block = None |
|
|
| if partial_block is not None: |
| append_completed(partial_block, total_sublayers) |
| return self.norm(read_residual(total_sublayers)) |
|
|
| def forward_hidden_cached( |
| self, |
| input_ids: torch.Tensor, |
| cache: YunmoNextCUDAState | None = None, |
| position_ids: torch.Tensor | None = None, |
| *, |
| layer_trace: list[torch.Tensor] | None = None, |
| routing_trace: list[torch.Tensor] | None = None, |
| ) -> tuple[torch.Tensor, YunmoNextCUDAState]: |
| """Exact single-document chunked prefill for the adopted model.""" |
|
|
| if self.training: |
| raise RuntimeError("cached inference requires model.eval()") |
| return self._forward_hidden_cached_from_embeddings( |
| self.embed_tokens(input_ids), |
| cache, |
| position_ids, |
| layer_trace=layer_trace, |
| routing_trace=routing_trace, |
| ) |
|
|
| def _forward_hidden_cached_from_embeddings( |
| self, |
| embeddings: torch.Tensor, |
| cache: YunmoNextCUDAState | None = None, |
| position_ids: torch.Tensor | None = None, |
| *, |
| layer_trace: list[torch.Tensor] | None = None, |
| routing_trace: list[torch.Tensor] | None = None, |
| ) -> tuple[torch.Tensor, YunmoNextCUDAState]: |
| if self.training: |
| raise RuntimeError("cached inference requires model.eval()") |
| if embeddings.ndim != 3: |
| raise ValueError( |
| "embeddings must have shape [batch, sequence, hidden]" |
| ) |
| batch, sequence, hidden_size = embeddings.shape |
| if hidden_size != self.config.hidden_size: |
| raise ValueError("embedding hidden size disagrees with model config") |
| if cache is None: |
| try: |
| from fla.models.utils import Cache |
| except ImportError as exc: |
| raise ImportError( |
| "cached recurrent inference requires FLA" |
| ) from exc |
| cache = YunmoNextCUDAState( |
| fla_cache=Cache(), |
| attention={}, |
| seen_tokens=0, |
| ) |
| if position_ids is None: |
| position_ids = torch.arange( |
| cache.seen_tokens, |
| cache.seen_tokens + sequence, |
| device=embeddings.device, |
| ).expand(batch, -1) |
| elif position_ids.shape != (batch, sequence): |
| raise ValueError("cached position IDs must match embeddings") |
|
|
| hidden_states = self._forward_hidden_two_phase_attention_residual_cached( |
| embeddings, |
| position_ids, |
| cache, |
| layer_trace, |
| routing_trace, |
| ) |
| cache.seen_tokens += sequence |
| return hidden_states, cache |
|
|
| def forward_cached( |
| self, |
| input_ids: torch.Tensor, |
| cache: YunmoNextCUDAState | None = None, |
| position_ids: torch.Tensor | None = None, |
| *, |
| materialize_logits: bool = True, |
| ) -> tuple[YunmoNextOutput, YunmoNextCUDAState]: |
| hidden_states, cache = self.forward_hidden_cached( |
| input_ids, |
| cache, |
| position_ids, |
| ) |
| logits = ( |
| F.linear(hidden_states, self.embed_tokens.weight) |
| if materialize_logits |
| else None |
| ) |
| return YunmoNextOutput( |
| logits=logits, |
| hidden_states=hidden_states, |
| ), cache |
|
|
| def _forward_cached_from_embeddings( |
| self, |
| embeddings: torch.Tensor, |
| cache: YunmoNextCUDAState | None = None, |
| position_ids: torch.Tensor | None = None, |
| *, |
| materialize_logits: bool = True, |
| ) -> tuple[YunmoNextOutput, YunmoNextCUDAState]: |
| hidden_states, cache = self._forward_hidden_cached_from_embeddings( |
| embeddings, |
| cache, |
| position_ids, |
| ) |
| logits = ( |
| F.linear(hidden_states, self.embed_tokens.weight) |
| if materialize_logits |
| else None |
| ) |
| return YunmoNextOutput( |
| logits=logits, |
| hidden_states=hidden_states, |
| ), cache |
|
|
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| labels: torch.Tensor | None = None, |
| position_ids: torch.Tensor | None = None, |
| cu_seqlens: torch.Tensor | None = None, |
| materialize_logits: bool = False, |
| ) -> YunmoNextOutput: |
| return self._forward_from_embeddings( |
| self.embed_tokens(input_ids), |
| labels=labels, |
| position_ids=position_ids, |
| cu_seqlens=cu_seqlens, |
| materialize_logits=materialize_logits, |
| ) |
|
|
| def _forward_from_embeddings( |
| self, |
| embeddings: torch.Tensor, |
| labels: torch.Tensor | None = None, |
| position_ids: torch.Tensor | None = None, |
| cu_seqlens: torch.Tensor | None = None, |
| materialize_logits: bool = False, |
| ) -> YunmoNextOutput: |
| hidden_states = self._forward_hidden_from_embeddings( |
| embeddings, |
| position_ids, |
| cu_seqlens, |
| ) |
|
|
| loss = None |
| logits = None |
| if ( |
| labels is not None |
| and self.training_engine.cross_entropy_backend == "liger" |
| ): |
| try: |
| from liger_kernel.transformers import ( |
| LigerFusedLinearCrossEntropyLoss, |
| ) |
| except ImportError as exc: |
| raise ImportError( |
| "the configured cross-entropy backend requires liger-kernel" |
| ) from exc |
| loss_fn = LigerFusedLinearCrossEntropyLoss( |
| accum_dtype=torch.float32 |
| ) |
| loss = loss_fn( |
| self.embed_tokens.weight, |
| hidden_states.reshape(-1, self.config.hidden_size), |
| labels.reshape(-1), |
| ) |
| elif labels is not None or materialize_logits: |
| logits = F.linear(hidden_states, self.embed_tokens.weight) |
| if labels is not None: |
| loss = F.cross_entropy( |
| logits.float().reshape(-1, self.config.vocab_size), |
| labels.reshape(-1), |
| ) |
|
|
| return YunmoNextOutput( |
| loss=loss, |
| logits=logits, |
| hidden_states=hidden_states, |
| ) |
|
|
