diff --git a/docs/source/en/model_doc/rwkv.md b/docs/source/en/model_doc/rwkv.md index 0c26b717f5fa..81c2f0c1329e 100644 --- a/docs/source/en/model_doc/rwkv.md +++ b/docs/source/en/model_doc/rwkv.md @@ -15,7 +15,7 @@ specific language governing permissions and limitations under the License. ## Overview -RWKV-7 is a recurrent language model whose TimeMix update uses a diagonal-plus-low-rank state transition. This +RWKV-7 is a recurrent language model whose linear-attention update uses a diagonal-plus-low-rank state transition. This implementation replaces the former RWKV-4 model behind the existing `rwkv` model identity. RWKV-4 checkpoints are not compatible and are rejected explicitly. @@ -24,8 +24,8 @@ serialization. All product computation is delegated to the public [FlashRWKV2](h operator API. Training follows `RWKV-LM/RWKV-v7/train_temp`; inference follows Albatross. There is no CPU, PyTorch or FLA product fallback. -Inference calls FlashRWKV2 at model-semantic fusion boundaries: each TimeMix layer uses PostNorm+TokenShift, WKV -Prepare, WKV7 and Readout; each ChannelMix layer uses one complete ChannelMix operator. Transformers does not select +Inference calls FlashRWKV2 at model-semantic fusion boundaries: each linear-attention layer uses PostNorm+TokenShift, +WKV Prepare, WKV7 and Readout; each MLP uses one complete fused operator. Transformers does not select sparse/dense kernels or invoke standalone projection, activation, LN, Res, TokenShift, VRes or gate helpers. The current canonical contract uses: @@ -34,7 +34,7 @@ The current canonical contract uses: - BF16 CUDA tensors and sequence lengths divisible by 16 for pretraining; - Albatross's mixed BF16 embedding / FP16 model layout for inference; - FP32 recurrent WKV state; -- one active unmerged vanilla LoRA adapter on the TimeMix `receptance`, `key`, `value`, and `output` projections; +- one active unmerged vanilla LoRA adapter on the linear-attention `r_proj`, `k_proj`, `v_proj`, and `o_proj` modules; multiple adapters, LoRA variants, LoRA bias, and per-sample mixed-adapter batches must be merged first; - equal-length batches for training and inference. An attention mask may be omitted or may be a two-dimensional, batch-matched, all-ones tensor whose length covers the current input. Padding and ragged batches fail immediately; @@ -78,7 +78,7 @@ outputs = model.generate(**inputs, max_new_tokens=32) `prepare_for_inference()` is an in-place conversion, not a temporary execution mode. Call it after loading or changing weights and after attaching adapters. It changes parameter dtypes and device placement, creates non-persistent -transposed runtime weights, and moves the serializable ChannelMix down-projection weights to CPU so a second full GPU +transposed runtime weights, and moves the serializable MLP `down_proj` weights to CPU so a second full GPU copy is not retained. Repeating the call is safe. Saving still serializes the canonical weights, and loading the saved checkpoint starts without runtime layouts. @@ -124,19 +124,23 @@ Canonical BlinkDL `.pth` checkpoints can be converted with: --context-length 10240 ``` -The converter preserves canonical tensor names under the single standard `model.` base-model prefix and writes -Safetensors without a per-tensor compatibility table. It builds a standard fast `tokenizer.json` from the pinned +The converter treats BlinkDL names as an input format and writes the native Transformers decoder layout. For example, +`blocks.0.att.receptance.weight` becomes `model.layers.0.linear_attn.r_proj.weight`, +`blocks.0.ffn.value.weight` becomes `model.layers.0.mlp.down_proj.weight`, and `head.weight` becomes +`lm_head.weight`. The converted Safetensors checkpoint therefore works with standard module discovery, adapter target +selection, and framework tooling without preserving a second set of runtime aliases. The converter also builds a +standard fast `tokenizer.json` from the pinned `rwkv-rs/rwkv7-g1-st` `rwkv_vocab_v20230424.json` artifact and writes the native RWKV chat template. Pass `--rwkv-vocab-json` to use a hash-verified local copy of that JSON artifact or `--chat-template` to select the template file. The tokenizer uses the RWKV World byte-level greedy longest-match algorithm. BOS and EOS share token ID 0; padding and unknown tokens are intentionally undefined. -For a new model, the four TimeMix low-rank dimensions are derived from `hidden_size` with the exact `train_temp` +For a new model, the four linear-attention low-rank dimensions are derived from `hidden_size` with the exact `train_temp` formulas. Converted checkpoints instead record the dimensions found in `w1/a1/v1/g1` and preserve them verbatim; loading a checkpoint never recomputes or replaces its serialized low-rank dimensions. Randomly initialized models also use the final effective parameter initialization from `train_temp`, including SmallInitEmb, the -vocabulary-dependent orthogonal LM head, layer-scaled GroupNorm weights, and the depth-dependent TimeMix and -ChannelMix parameters. +vocabulary-dependent orthogonal LM head, layer-scaled GroupNorm weights, and the depth-dependent linear-attention and +MLP parameters. ## RwkvTokenizerFast @@ -146,9 +150,9 @@ ChannelMix parameters. [[autodoc]] RwkvConfig -## RwkvTimeMix +## RwkvLinearAttention -[[autodoc]] RwkvTimeMix +[[autodoc]] RwkvLinearAttention - forward ## RwkvCache diff --git a/src/transformers/models/rwkv/convert_rwkv_checkpoint_to_hf.py b/src/transformers/models/rwkv/convert_rwkv_checkpoint_to_hf.py deleted file mode 100644 index 61ce995615a6..000000000000 --- a/src/transformers/models/rwkv/convert_rwkv_checkpoint_to_hf.py +++ /dev/null @@ -1,208 +0,0 @@ -# Copyright 2023 The HuggingFace Inc. team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Convert a RWKV checkpoint from BlinkDL to the Hugging Face format.""" - -import argparse -import gc -import json -import os -import re - -import torch -from huggingface_hub import hf_hub_download, split_torch_state_dict_into_shards - -from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizerFast, RwkvConfig -from transformers.modeling_utils import WEIGHTS_INDEX_NAME - - -NUM_HIDDEN_LAYERS_MAPPING = { - "169M": 12, - "430M": 24, - "1B5": 24, - "3B": 32, - "7B": 32, - "14B": 40, -} - -HIDDEN_SIZE_MAPPING = { - "169M": 768, - "430M": 1024, - "1B5": 2048, - "3B": 2560, - "7B": 4096, - "14B": 5120, -} - - -def convert_state_dict(state_dict): - state_dict_keys = list(state_dict.keys()) - for name in state_dict_keys: - weight = state_dict.pop(name) - # emb -> embedding - if name.startswith("emb."): - name = name.replace("emb.", "embeddings.") - # ln_0 -> pre_ln (only present at block 0) - if name.startswith("blocks.0.ln0"): - name = name.replace("blocks.0.ln0", "blocks.0.pre_ln") - # att -> attention - name = re.sub(r"blocks\.(\d+)\.att", r"blocks.\1.attention", name) - # ffn -> feed_forward - name = re.sub(r"blocks\.(\d+)\.ffn", r"blocks.\1.feed_forward", name) - # time_mix_k -> time_mix_key and reshape - if name.endswith(".time_mix_k"): - name = name.replace(".time_mix_k", ".time_mix_key") - # time_mix_v -> time_mix_value and reshape - if name.endswith(".time_mix_v"): - name = name.replace(".time_mix_v", ".time_mix_value") - # time_mix_r -> time_mix_key and reshape - if name.endswith(".time_mix_r"): - name = name.replace(".time_mix_r", ".time_mix_receptance") - - if name != "head.weight": - name = "rwkv." + name - - state_dict[name] = weight - return state_dict - - -def convert_rmkv_checkpoint_to_hf_format( - repo_id, checkpoint_file, output_dir, size=None, tokenizer_file=None, push_to_hub=False, model_name=None -): - # 1. If possible, build the tokenizer. - if tokenizer_file is None: - print("No `--tokenizer_file` provided, we will use the default tokenizer.") - vocab_size = 50277 - tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b") - else: - tokenizer = PreTrainedTokenizerFast(tokenizer_file=tokenizer_file) - vocab_size = len(tokenizer) - tokenizer.save_pretrained(output_dir) - - # 2. Build the config - possible_sizes = list(NUM_HIDDEN_LAYERS_MAPPING.keys()) - if size is None: - # Try to infer size from the checkpoint name - for candidate in possible_sizes: - if candidate in checkpoint_file: - size = candidate - break - if size is None: - raise ValueError("Could not infer the size, please provide it with the `--size` argument.") - if size not in possible_sizes: - raise ValueError(f"`size` should be one of {possible_sizes}, got {size}.") - - config = RwkvConfig( - vocab_size=vocab_size, - num_hidden_layers=NUM_HIDDEN_LAYERS_MAPPING[size], - hidden_size=HIDDEN_SIZE_MAPPING[size], - ) - config.save_pretrained(output_dir) - - # 3. Download model file then convert state_dict - model_file = hf_hub_download(repo_id, checkpoint_file) - state_dict = torch.load(model_file, map_location="cpu", weights_only=True) - state_dict = convert_state_dict(state_dict) - - # 4. Split in shards and save - state_dict_split = split_torch_state_dict_into_shards(state_dict) - shards = index = None - for tensors in state_dict_split.filename_to_tensors.values(): - shards = {tensor: state_dict[tensor] for tensor in tensors} - if state_dict_split.is_sharded: - index = { - "metadata": state_dict_split.metadata, - "weight_map": state_dict_split.tensor_to_filename, - } - - for shard_file, shard in shards.items(): - torch.save(shard, os.path.join(output_dir, shard_file)) - - if index is not None: - save_index_file = os.path.join(output_dir, WEIGHTS_INDEX_NAME) - # Save the index as well - with open(save_index_file, "w", encoding="utf-8") as f: - content = json.dumps(index, indent=2, sort_keys=True) + "\n" - f.write(content) - - # 5. Clean up shards (for some reason the file PyTorch saves take the same space as the whole state_dict - print( - "Cleaning up shards. This may error with an OOM error, it this is the case don't worry you still have converted the model." - ) - shard_files = list(shards.keys()) - - del state_dict - del shards - gc.collect() - - for shard_file in shard_files: - state_dict = torch.load(os.path.join(output_dir, shard_file), weights_only=True) - torch.save({k: v.cpu().clone() for k, v in state_dict.items()}, os.path.join(output_dir, shard_file)) - - del state_dict - gc.collect() - - if push_to_hub: - if model_name is None: - raise ValueError("Please provide a `model_name` to push the model to the Hub.") - model = AutoModelForCausalLM.from_pretrained(output_dir) - model.push_to_hub(model_name, max_shard_size="2GB") - tokenizer.push_to_hub(model_name) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - # Required parameters - parser.add_argument( - "--repo_id", default=None, type=str, required=True, help="Repo ID from which to pull the checkpoint." - ) - parser.add_argument( - "--checkpoint_file", default=None, type=str, required=True, help="Name of the checkpoint file in the repo." - ) - parser.add_argument( - "--output_dir", default=None, type=str, required=True, help="Where to save the converted model." - ) - parser.add_argument( - "--tokenizer_file", - default=None, - type=str, - help="Path to the tokenizer file to use (if not provided, only the model is converted).", - ) - parser.add_argument( - "--size", - default=None, - type=str, - help="Size of the model. Will be inferred from the `checkpoint_file` if not passed.", - ) - parser.add_argument( - "--push_to_hub", - action="store_true", - help="Push to the Hub the converted model.", - ) - parser.add_argument( - "--model_name", - default=None, - type=str, - help="Name of the pushed model on the Hub, including the username / organization.", - ) - - args = parser.parse_args() - convert_rmkv_checkpoint_to_hf_format( - args.repo_id, - args.checkpoint_file, - args.output_dir, - size=args.size, - tokenizer_file=args.tokenizer_file, - push_to_hub=args.push_to_hub, - model_name=args.model_name, - ) diff --git a/src/transformers/models/rwkv/modeling_rwkv.py b/src/transformers/models/rwkv/modeling_rwkv.py index 0be68eda6a78..63177920166c 100644 --- a/src/transformers/models/rwkv/modeling_rwkv.py +++ b/src/transformers/models/rwkv/modeling_rwkv.py @@ -91,7 +91,7 @@ def _stateful_training_metadata( return sequence_chunk_offsets, chunk_token_starts, chunk_token_ends -def _infer_tmix_projection_spec( +def _infer_linear_attention_projection_spec( projection: nn.Module, ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, float]: """Resolve one bias-free Linear and an optional active vanilla LoRA adapter.""" @@ -121,7 +121,8 @@ def _infer_tmix_projection_spec( missing = [name for name in required if not hasattr(projection, name)] if missing: raise RuntimeError( - "RWKV-7 FlashRWKV2 inference only supports vanilla PEFT LoRA wrappers around TimeMix projections; " + "RWKV-7 FlashRWKV2 inference only supports vanilla PEFT LoRA wrappers around linear-attention " + "projections; " f"{type(projection).__name__} is missing {missing}." ) if getattr(projection, "fan_in_fan_out", False): @@ -187,9 +188,9 @@ class RwkvTrainingState: WKV accumulation is always FP32. """ - time_mix_shift: torch.Tensor - wkv: torch.Tensor - channel_mix_shift: torch.Tensor + attention_shift: torch.Tensor + recurrent_state: torch.Tensor + mlp_shift: torch.Tensor @classmethod def zeros( @@ -213,9 +214,9 @@ def zeros( config.head_size, ) return cls( - time_mix_shift=torch.zeros(shift_shape, device=device, dtype=dtype), - wkv=torch.zeros(wkv_shape, device=device, dtype=torch.float32), - channel_mix_shift=torch.zeros(shift_shape, device=device, dtype=dtype), + attention_shift=torch.zeros(shift_shape, device=device, dtype=dtype), + recurrent_state=torch.zeros(wkv_shape, device=device, dtype=torch.float32), + mlp_shift=torch.zeros(shift_shape, device=device, dtype=dtype), ) def validate( @@ -236,9 +237,9 @@ def validate( config.head_size, ) fields = ( - ("time_mix_shift", self.time_mix_shift, expected_shift_shape, dtype), - ("wkv", self.wkv, expected_wkv_shape, torch.float32), - ("channel_mix_shift", self.channel_mix_shift, expected_shift_shape, dtype), + ("attention_shift", self.attention_shift, expected_shift_shape, dtype), + ("recurrent_state", self.recurrent_state, expected_wkv_shape, torch.float32), + ("mlp_shift", self.mlp_shift, expected_shift_shape, dtype), ) for name, value, shape, expected_dtype in fields: if not isinstance(value, torch.Tensor): @@ -261,20 +262,20 @@ def clone_detach(self) -> RwkvTrainingState: return type(self)(*(value.clone().detach() for value in self.tensors())) def tensors(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - return self.time_mix_shift, self.wkv, self.channel_mix_shift + return self.attention_shift, self.recurrent_state, self.mlp_shift def reset_( self, batch_indices: torch.Tensor | list[int] | tuple[int, ...] | None = None, *, - time_mix: bool = True, - wkv: bool = True, - channel_mix: bool = True, + attention: bool = True, + recurrent: bool = True, + mlp: bool = True, ) -> RwkvTrainingState: """Zero selected batch rows in place, or all rows when indices are omitted.""" - if not any((time_mix, wkv, channel_mix)): + if not any((attention, recurrent, mlp)): return self - selected = (time_mix, wkv, channel_mix) + selected = (attention, recurrent, mlp) with torch.no_grad(): for enabled, value in zip(selected, self.tensors(), strict=True): if not enabled: @@ -290,16 +291,16 @@ def reset( self, batch_indices: torch.Tensor | list[int] | tuple[int, ...] | None = None, *, - time_mix: bool = True, - wkv: bool = True, - channel_mix: bool = True, + attention: bool = True, + recurrent: bool = True, + mlp: bool = True, ) -> RwkvTrainingState: """Return a clone with selected batch rows reset to zero.""" return self.clone().reset_( batch_indices, - time_mix=time_mix, - wkv=wkv, - channel_mix=channel_mix, + attention=attention, + recurrent=recurrent, + mlp=mlp, ) @@ -484,8 +485,8 @@ def reset_parameters(self) -> None: init.orthogonal_(self.weight, gain=gain) -class RwkvTimeMix(nn.Module): - """Canonical RWKV-7 TimeMix component using FlashRWKV2's public training and inference APIs.""" +class RwkvLinearAttention(nn.Module): + """RWKV-7 linear attention using FlashRWKV2's public training and inference APIs.""" def __init__(self, config: RwkvConfig, layer_idx: int): super().__init__() @@ -512,11 +513,11 @@ def __init__(self, config: RwkvConfig, layer_idx: int): self.k_k = nn.Parameter(torch.empty(1, 1, channels)) self.k_a = nn.Parameter(torch.empty(1, 1, channels)) self.r_k = nn.Parameter(torch.empty(heads, config.head_size)) - self.receptance = nn.Linear(channels, channels, bias=False) - self.key = nn.Linear(channels, channels, bias=False) - self.value = nn.Linear(channels, channels, bias=False) - self.output = nn.Linear(channels, channels, bias=False) - self.ln_x = nn.GroupNorm(heads, channels, eps=config.group_norm_epsilon) + self.r_proj = nn.Linear(channels, channels, bias=False) + self.k_proj = nn.Linear(channels, channels, bias=False) + self.v_proj = nn.Linear(channels, channels, bias=False) + self.o_proj = nn.Linear(channels, channels, bias=False) + self.g_norm = nn.GroupNorm(heads, channels, eps=config.group_norm_epsilon) for name in ("w1", "w2", "a1", "a2", "v1", "v2", "g1", "g2"): self.register_buffer(f"_{name}_original", None, persistent=False) self.register_load_state_dict_post_hook(self._clear_inference_layouts) @@ -587,25 +588,25 @@ def reset_parameters(self) -> None: else 1 ) init.orthogonal_(parameter, gain=gain * 0.1) - init.orthogonal_(self.receptance.weight, gain=1.0) - init.orthogonal_(self.key.weight, gain=0.1) - init.orthogonal_(self.value.weight, gain=1.0) - init.zeros_(self.output.weight) + init.orthogonal_(self.r_proj.weight, gain=1.0) + init.orthogonal_(self.k_proj.weight, gain=0.1) + init.orthogonal_(self.v_proj.weight, gain=1.0) + init.zeros_(self.o_proj.weight) layer_scale = (self.layer_idx + 1) / self.config.num_hidden_layers - init.constant_(self.ln_x.weight, layer_scale**0.7) - init.zeros_(self.ln_x.bias) + init.constant_(self.g_norm.weight, layer_scale**0.7) + init.zeros_(self.g_norm.bias) def _training_projections(self, flash, mixed: tuple[torch.Tensor, ...], v_first: torch.Tensor | None): xr, xw, xk, xv, xa, xg = mixed - receptance = self.receptance(xr) + receptance = self.r_proj(xr) decay_logits = self.w0 + torch.tanh(xw @ self.w1) @ self.w2 - key = self.key(xk) - value = self.value(xv) + key = self.k_proj(xk) + value = self.v_proj(xv) if self.layer_idx == 0: v_first = value else: if v_first is None: - raise ValueError("`v_first` must be supplied to RWKV-7 TimeMix layers after layer 0.") + raise ValueError("`v_first` must be supplied to RWKV-7 linear-attention layers after layer 0.") value = flash.pretrain_tmix_vres_gate_bf16( value.contiguous(), v_first.contiguous(), @@ -639,11 +640,11 @@ def _finish_training_output( key, value.contiguous(), self.r_k.contiguous(), - self.ln_x.weight.contiguous(), - self.ln_x.bias.contiguous(), + self.g_norm.weight.contiguous(), + self.g_norm.bias.contiguous(), gate, ) - return self.output(output) + return self.o_proj(output) def _training_forward(self, hidden_states: torch.Tensor, v_first: torch.Tensor | None): flash = _load_flash_rwkv2("training", hidden_states) @@ -777,14 +778,16 @@ def _inference_forward( eps=layer_norm.eps, validated_metadata=ticket, ) - receptance_weight, receptance_lora_a, receptance_lora_b, receptance_lora_scale = _infer_tmix_projection_spec( - self.receptance + receptance_weight, receptance_lora_a, receptance_lora_b, receptance_lora_scale = ( + _infer_linear_attention_projection_spec(self.r_proj) + ) + key_weight, key_lora_a, key_lora_b, key_lora_scale = _infer_linear_attention_projection_spec(self.k_proj) + value_weight, value_lora_a, value_lora_b, value_lora_scale = _infer_linear_attention_projection_spec( + self.v_proj ) - key_weight, key_lora_a, key_lora_b, key_lora_scale = _infer_tmix_projection_spec(self.key) - value_weight, value_lora_a, value_lora_b, value_lora_scale = _infer_tmix_projection_spec(self.value) w1, w2, a1, a2, v1, v2, g1, g2 = self._inference_low_rank_layouts() if self.layer_idx != 0 and v_first is None: - raise ValueError("`v_first` must be supplied to RWKV-7 TimeMix layers after layer 0.") + raise ValueError("`v_first` must be supplied to RWKV-7 linear-attention layers after layer 0.") ( receptance, decay_delta, @@ -854,15 +857,17 @@ def _inference_forward( max_seqlen=sequence_length, validated_metadata=ticket, ).view(-1, channels) - output_weight, output_lora_a, output_lora_b, output_lora_scale = _infer_tmix_projection_spec(self.output) + output_weight, output_lora_a, output_lora_b, output_lora_scale = _infer_linear_attention_projection_spec( + self.o_proj + ) output = flash.infer_tmix_readout_forward_varlen( output, receptance, key, value, self.r_k.reshape(-1).contiguous(), - self.ln_x.weight.contiguous(), - self.ln_x.bias.contiguous(), + self.g_norm.weight.contiguous(), + self.g_norm.bias.contiguous(), gate, output_weight, output_lora_a=output_lora_a, @@ -910,41 +915,42 @@ def forward( "RWKV-7 inference requires an RwkvCache, including when the caller discards the final cache." ) raise RuntimeError( - "RWKV-7 inference must run TimeMix through its owning block so FlashRWKV2 can fuse residual, " + "RWKV-7 inference must run linear attention through its owning decoder layer so FlashRWKV2 can fuse " + "residual, " "LayerNorm and TokenShift." ) -class RwkvChannelMix(nn.Module): +class RwkvMLP(nn.Module): def __init__(self, config: RwkvConfig, layer_idx: int): super().__init__() self.config = config self.layer_idx = layer_idx self.x_k = nn.Parameter(torch.empty(1, 1, config.hidden_size)) - self.key = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) - self.value = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) - self.register_buffer("_value_runtime", None, persistent=False) + self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) + self.register_buffer("_down_proj_runtime", None, persistent=False) self.register_load_state_dict_post_hook(self._clear_inference_layout) def _clear_inference_layout(self, *args) -> None: - self._value_runtime = None + self._down_proj_runtime = None def _apply(self, fn, recurse=True): self._clear_inference_layout() return super()._apply(fn, recurse=recurse) def prepare_for_inference(self) -> None: - runtime_device = self.key.weight.device - if runtime_device.type != "cuda" or self.key.weight.dtype != torch.float16: + runtime_device = self.up_proj.weight.device + if runtime_device.type != "cuda" or self.up_proj.weight.dtype != torch.float16: raise RuntimeError( - "RWKV-7 Albatross ChannelMix layout requires CUDA float16 runtime weights; " - f"got dtype={self.key.weight.dtype}, device={runtime_device}." + "RWKV-7 Albatross MLP layout requires CUDA float16 runtime weights; " + f"got dtype={self.up_proj.weight.dtype}, device={runtime_device}." ) - self._value_runtime = self.value.weight.to(device=runtime_device, dtype=torch.float16).T.contiguous() + self._down_proj_runtime = self.down_proj.weight.to(device=runtime_device, dtype=torch.float16).T.contiguous() # Albatross replaces the canonical FFN-down layout during inference. Keep the serializable parameter on CPU # instead of retaining a second 4 GiB GPU copy for a 7.2B model; the non-persistent runtime layout is the only # one consumed after this explicit inference preparation step. - self.value.weight.data = self.value.weight.data.cpu() + self.down_proj.weight.data = self.down_proj.weight.data.cpu() def reset_parameters(self) -> None: with torch.no_grad(): @@ -952,8 +958,8 @@ def reset_parameters(self) -> None: ratio_1_to_almost0 = 1.0 - self.layer_idx / self.config.num_hidden_layers ddd = torch.arange(channels, dtype=torch.float32, device=self.x_k.device).view(1, 1, -1) / channels init.copy_(self.x_k, 1.0 - ddd.pow(ratio_1_to_almost0**4)) - init.orthogonal_(self.key.weight, gain=1.0) - init.zeros_(self.value.weight) + init.orthogonal_(self.up_proj.weight, gain=1.0) + init.zeros_(self.down_proj.weight) def forward( self, @@ -967,27 +973,27 @@ def forward( if training_shift_state is not None: if hidden_states.dtype != torch.bfloat16: raise RuntimeError( - f"RWKV-7 stateful ChannelMix requires bfloat16 activations; got {hidden_states.dtype}." + f"RWKV-7 stateful MLP requires bfloat16 activations; got {hidden_states.dtype}." ) flash = _load_flash_rwkv2("stateful training", hidden_states) return flash.statetune_cmix_bf16( hidden_states.contiguous(), training_shift_state.contiguous(), self.x_k.reshape(-1).contiguous(), - self.key.weight.contiguous(), - self.value.weight.contiguous(), + self.up_proj.weight.contiguous(), + self.down_proj.weight.contiguous(), ) flash = _load_flash_rwkv2("training", hidden_states) return flash.pretrain_cmix_bf16( hidden_states.contiguous(), self.x_k.reshape(-1).contiguous(), - self.key.weight.contiguous(), - self.value.weight.contiguous(), + self.up_proj.weight.contiguous(), + self.down_proj.weight.contiguous(), ) if past_key_values is None: raise ValueError("RWKV-7 inference requires an RwkvCache.") raise RuntimeError( - "RWKV-7 inference must run ChannelMix through its owning block so FlashRWKV2 can fuse residual, " + "RWKV-7 inference must run the MLP through its owning decoder layer so FlashRWKV2 can fuse residual, " "LayerNorm, TokenShift and the complete FFN." ) @@ -998,11 +1004,11 @@ def inference_forward_with_postnorm( layer_norm: nn.LayerNorm, past_key_values: RwkvCache, ) -> tuple[torch.Tensor, torch.Tensor]: - """Run the complete FlashRWKV2 ChannelMix inference island.""" + """Run the complete FlashRWKV2 MLP inference island.""" flash = _load_flash_rwkv2("inference", hidden_states) - if self._value_runtime is None: + if self._down_proj_runtime is None: raise RuntimeError( - "RWKV-7 Albatross ChannelMix layout is not prepared; call `model.prepare_for_inference()` " + "RWKV-7 Albatross MLP layout is not prepared; call `model.prepare_for_inference()` " "after loading or modifying weights." ) batch_size, sequence_length, channels = hidden_states.shape @@ -1017,8 +1023,8 @@ def inference_forward_with_postnorm( layer_norm.weight.contiguous(), layer_norm.bias.contiguous(), self.x_k.reshape(-1).contiguous(), - self.key.weight.contiguous(), - self._value_runtime, + self.up_proj.weight.contiguous(), + self._down_proj_runtime, shift_state_pool=ffn_shift, cu_seqlens=cu_seqlens, state_indices=state_indices, @@ -1030,17 +1036,15 @@ def inference_forward_with_postnorm( return summed.view_as(hidden_states), output.view_as(hidden_states) -class RwkvBlock(GradientCheckpointingLayer): +class RwkvDecoderLayer(GradientCheckpointingLayer): def __init__(self, config: RwkvConfig, layer_idx: int): super().__init__() self.layer_idx = layer_idx self.hidden_state_boundary = nn.Identity() - if layer_idx == 0: - self.ln0 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) - self.ln1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) - self.ln2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) - self.att = RwkvTimeMix(config, layer_idx) - self.ffn = RwkvChannelMix(config, layer_idx) + self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) + self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) + self.linear_attn = RwkvLinearAttention(config, layer_idx) + self.mlp = RwkvMLP(config, layer_idx) def forward( self, @@ -1054,25 +1058,25 @@ def forward( if not self.training: if past_key_values is None: raise ValueError("RWKV-7 inference requires an RwkvCache.") - raise RuntimeError("RWKV-7 inference blocks require an explicit residual tensor.") - att_result = self.att( - self.ln1(hidden_states), + raise RuntimeError("RWKV-7 inference decoder layers require an explicit residual tensor.") + att_result = self.linear_attn( + self.input_layernorm(hidden_states), v_first=v_first, past_key_values=past_key_values, attention_mask=attention_mask, - training_shift_state=None if training_state is None else training_state.time_mix_shift[self.layer_idx], - training_wkv_state=None if training_state is None else training_state.wkv[self.layer_idx], + training_shift_state=None if training_state is None else training_state.attention_shift[self.layer_idx], + training_wkv_state=None if training_state is None else training_state.recurrent_state[self.layer_idx], ) if training_state is None: output, v_first = att_result else: output, v_first, next_att_shift, next_wkv = att_result hidden_states = hidden_states + output - ffn_result = self.ffn( - self.ln2(hidden_states), + ffn_result = self.mlp( + self.post_attention_layernorm(hidden_states), past_key_values=past_key_values, attention_mask=attention_mask, - training_shift_state=None if training_state is None else training_state.channel_mix_shift[self.layer_idx], + training_shift_state=None if training_state is None else training_state.mlp_shift[self.layer_idx], ) if training_state is None: hidden_states = hidden_states + ffn_result @@ -1090,25 +1094,25 @@ def inference_forward( v_first: torch.Tensor | None, past_key_values: RwkvCache, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Run one inference block while carrying FlashRWKV2's residual between fusion islands.""" - block_input, output, v_first = self.att.inference_forward_with_postnorm( - hidden_states, residual, self.ln1, v_first, past_key_values + """Run one decoder layer while carrying FlashRWKV2's residual between fusion islands.""" + layer_input, output, v_first = self.linear_attn.inference_forward_with_postnorm( + hidden_states, residual, self.input_layernorm, v_first, past_key_values ) - block_input = self.hidden_state_boundary(block_input) - hidden_states, residual = self.ffn.inference_forward_with_postnorm( - block_input, output, self.ln2, past_key_values + layer_input = self.hidden_state_boundary(layer_input) + hidden_states, residual = self.mlp.inference_forward_with_postnorm( + layer_input, output, self.post_attention_layernorm, past_key_values ) layer = past_key_values.layers[self.layer_idx] if isinstance(layer, RwkvDynamicCacheLayer): layer.mark_updated(hidden_states.shape[1]) - return hidden_states, residual, v_first, block_input + return hidden_states, residual, v_first, layer_input @auto_docstring class RwkvPreTrainedModel(PreTrainedModel): config_class = RwkvConfig base_model_prefix = "model" - _no_split_modules = ["RwkvBlock"] + _no_split_modules = ["RwkvDecoderLayer"] _is_stateful = True _can_compile_fullgraph = False _can_record_outputs = { @@ -1124,7 +1128,7 @@ class RwkvPreTrainedModel(PreTrainedModel): def _init_weights(self, module): # These owning modules preserve Transformers' per-module loading markers, so from_pretrained never # reinitializes a complete model after loading a checkpoint. - if isinstance(module, RwkvEmbedding | RwkvLMHead | RwkvTimeMix | RwkvChannelMix): + if isinstance(module, RwkvEmbedding | RwkvLMHead | RwkvLinearAttention | RwkvMLP): module.reset_parameters() elif isinstance(module, nn.LayerNorm): module.reset_parameters() @@ -1138,31 +1142,33 @@ def __init__(self, config: RwkvConfig): f"RwkvModel only supports `architecture_version='rwkv7'`, got {config.architecture_version!r}." ) super().__init__(config) - self.emb = RwkvEmbedding(config.vocab_size, config.hidden_size) - self.blocks = nn.ModuleList([RwkvBlock(config, index) for index in range(config.num_hidden_layers)]) - self.ln_out = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) + self.embed_tokens = RwkvEmbedding(config.vocab_size, config.hidden_size) + self.embedding_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) + self.layers = nn.ModuleList([RwkvDecoderLayer(config, index) for index in range(config.num_hidden_layers)]) + self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon) self.hidden_state_boundary = nn.Identity() self.post_init() def reset_parameters(self) -> None: """Apply the final canonical train_temp initialization in model order.""" with torch.no_grad(): - init.uniform_(self.emb.weight, -1e-4, 1e-4) - for block in self.blocks: - for layer_norm in (getattr(block, "ln0", None), block.ln1, block.ln2): - if layer_norm is not None: - init.ones_(layer_norm.weight) - init.zeros_(layer_norm.bias) - block.att.reset_parameters() - block.ffn.reset_parameters() - init.ones_(self.ln_out.weight) - init.zeros_(self.ln_out.bias) + init.uniform_(self.embed_tokens.weight, -1e-4, 1e-4) + init.ones_(self.embedding_norm.weight) + init.zeros_(self.embedding_norm.bias) + for layer in self.layers: + for layer_norm in (layer.input_layernorm, layer.post_attention_layernorm): + init.ones_(layer_norm.weight) + init.zeros_(layer_norm.bias) + layer.linear_attn.reset_parameters() + layer.mlp.reset_parameters() + init.ones_(self.norm.weight) + init.zeros_(self.norm.bias) def get_input_embeddings(self): - return self.emb + return self.embed_tokens def set_input_embeddings(self, value): - self.emb = value + self.embed_tokens = value def _new_cache(self) -> RwkvCache: return RwkvCache(self.config) @@ -1170,12 +1176,12 @@ def _new_cache(self) -> RwkvCache: def prepare_for_inference(self): """Convert weights in-place to Albatross's mixed BF16-embedding/FP16-runtime layout.""" self.to(dtype=torch.float16) - self.emb.to(dtype=torch.bfloat16) + self.embed_tokens.to(dtype=torch.bfloat16) if not self.config.embedding_layer_norm_fused: - self.blocks[0].ln0.to(dtype=torch.bfloat16) - for block in self.blocks: - block.att.prepare_for_inference() - block.ffn.prepare_for_inference() + self.embedding_norm.to(dtype=torch.bfloat16) + for layer in self.layers: + layer.linear_attn.prepare_for_inference() + layer.mlp.prepare_for_inference() return self @merge_with_config_defaults @@ -1201,10 +1207,10 @@ def forward( if past_key_values is not None and not isinstance(past_key_values, RwkvCache): raise TypeError(f"RWKV-7 requires `RwkvCache`, got {type(past_key_values).__name__}.") if inputs_embeds is None: - inputs_embeds = self.emb(input_ids) + inputs_embeds = self.embed_tokens(input_ids) _validate_rwkv_attention_mask(attention_mask, inputs_embeds) if self.training: - if any(block.ffn._value_runtime is not None for block in self.blocks): + if any(layer.mlp._down_proj_runtime is not None for layer in self.layers): raise RuntimeError( "RWKV-7 is still using the in-place Albatross inference layout. To resume training, move the " "complete model to CUDA bfloat16 with `model.to(device='cuda', dtype=torch.bfloat16)`, then call " @@ -1221,7 +1227,7 @@ def forward( ) hidden_states = inputs_embeds if not self.config.embedding_layer_norm_fused: - hidden_states = self.blocks[0].ln0(hidden_states) + hidden_states = self.embedding_norm(hidden_states) cache = None else: if training_state is not None: @@ -1238,8 +1244,8 @@ def forward( flash = _load_flash_rwkv2("inference", inputs_embeds) hidden_states = flash.infer_embedding_ln0_forward_varlen( inputs_embeds.reshape(-1, self.config.hidden_size).contiguous(), - self.blocks[0].ln0.weight.contiguous(), - self.blocks[0].ln0.bias.contiguous(), + self.embedding_norm.weight.contiguous(), + self.embedding_norm.bias.contiguous(), eps=self.config.layer_norm_epsilon, ).view_as(inputs_embeds) @@ -1248,8 +1254,8 @@ def forward( next_wkv_states: list[torch.Tensor] = [] next_ffn_shifts: list[torch.Tensor] = [] if self.training: - for block in self.blocks: - block_result = block( + for layer in self.layers: + layer_result = layer( hidden_states, v_first, training_state, @@ -1257,26 +1263,26 @@ def forward( attention_mask=attention_mask, ) if training_state is None: - hidden_states, v_first = block_result + hidden_states, v_first = layer_result else: - hidden_states, v_first, next_att_shift, next_wkv, next_ffn_shift = block_result + hidden_states, v_first, next_att_shift, next_wkv, next_ffn_shift = layer_result next_att_shifts.append(next_att_shift) next_wkv_states.append(next_wkv) next_ffn_shifts.append(next_ffn_shift) - hidden_states = self.ln_out(hidden_states) + hidden_states = self.norm(hidden_states) else: flash = _load_flash_rwkv2("inference", hidden_states) if cache is None: raise RuntimeError("RWKV-7 inference cache initialization failed.") residual = torch.zeros_like(hidden_states) - for block in self.blocks: - hidden_states, residual, v_first, _ = block.inference_forward(hidden_states, residual, v_first, cache) + for layer in self.layers: + hidden_states, residual, v_first, _ = layer.inference_forward(hidden_states, residual, v_first, cache) batch_size, sequence_length, channels = hidden_states.shape hidden_states = flash.infer_post_norm_output_forward_varlen( hidden_states.reshape(-1, channels).contiguous(), residual.reshape(-1, channels).contiguous(), - self.ln_out.weight.contiguous(), - self.ln_out.bias.contiguous(), + self.norm.weight.contiguous(), + self.norm.bias.contiguous(), eps=self.config.layer_norm_epsilon, ).view(batch_size, sequence_length, channels) hidden_states = self.hidden_state_boundary(hidden_states) @@ -1285,9 +1291,9 @@ def forward( next_training_state = None if training_state is not None: next_training_state = RwkvTrainingState( - time_mix_shift=torch.stack(next_att_shifts), - wkv=torch.stack(next_wkv_states), - channel_mix_shift=torch.stack(next_ffn_shifts), + attention_shift=torch.stack(next_att_shifts), + recurrent_state=torch.stack(next_wkv_states), + mlp_shift=torch.stack(next_ffn_shifts), ) return RwkvModelOutput( last_hidden_state=hidden_states, @@ -1303,12 +1309,12 @@ class RwkvForCausalLM(RwkvPreTrainedModel, GenerationMixin): def __init__(self, config: RwkvConfig): super().__init__(config) self.model = RwkvModel(config) - self.head = RwkvLMHead(config) + self.lm_head = RwkvLMHead(config) self.post_init() - def reset_head_parameters(self) -> None: + def reset_lm_head_parameters(self) -> None: """Apply train_temp's vocabulary-dependent LM-head initialization.""" - self.head.reset_parameters() + self.lm_head.reset_parameters() def get_input_embeddings(self): return self.model.get_input_embeddings() @@ -1317,14 +1323,14 @@ def set_input_embeddings(self, value): self.model.set_input_embeddings(value) def get_output_embeddings(self): - return self.head + return self.lm_head def set_output_embeddings(self, value): - self.head = value + self.lm_head = value def prepare_for_inference(self): self.model.prepare_for_inference() - self.head.to(dtype=torch.float16) + self.lm_head.to(dtype=torch.float16) return self def prepare_inputs_for_generation( @@ -1391,20 +1397,20 @@ def forward( slice_indices = logits_to_keep selected_hidden_states = hidden_states[:, slice_indices, :] if self.training: - logits = self.head(selected_hidden_states) + logits = self.lm_head(selected_hidden_states) else: flash = _load_flash_rwkv2("inference", hidden_states) batch_size, sequence_length, channels = hidden_states.shape if isinstance(logits_to_keep, int) and logits_to_keep == 1: logits = flash.infer_head_linear_last_forward_varlen( selected_hidden_states.reshape(batch_size, channels).contiguous(), - self.head.weight.contiguous(), + self.lm_head.weight.contiguous(), tokens_count=sequence_length, ).view(batch_size, 1, self.config.vocab_size) else: selected_length = selected_hidden_states.shape[1] logits = flash.infer_head_linear_all_forward_varlen( - selected_hidden_states.reshape(-1, channels).contiguous(), self.head.weight.contiguous() + selected_hidden_states.reshape(-1, channels).contiguous(), self.lm_head.weight.contiguous() ).view(batch_size, selected_length, self.config.vocab_size) loss = None if labels is not None: @@ -1428,6 +1434,6 @@ def forward( "RwkvForCausalLM", "RwkvModel", "RwkvPreTrainedModel", - "RwkvTimeMix", + "RwkvLinearAttention", "RwkvTrainingState", ] diff --git a/temp/rwkv_pth2st.py b/temp/rwkv_pth2st.py index 17d202494293..6858914277d7 100755 --- a/temp/rwkv_pth2st.py +++ b/temp/rwkv_pth2st.py @@ -87,7 +87,7 @@ def infer_config(state: dict[str, torch.Tensor], context_length: int) -> RwkvCon raise ValueError("`blocks.0.ffn.key.weight` does not match the embedding hidden size.") if intermediate_size != 4 * hidden_size: raise ValueError( - "`blocks.0.ffn.key.weight` must use the canonical ChannelMix width 4 * hidden_size; " + "`blocks.0.ffn.key.weight` must use the canonical MLP width 4 * hidden_size; " f"got intermediate_size={intermediate_size} and hidden_size={hidden_size}." ) decay_rank = _infer_low_rank_dim(state, layer_ids, "w", hidden_size) @@ -111,10 +111,48 @@ def infer_config(state: dict[str, torch.Tensor], context_length: int) -> RwkvCon ) +def _transformers_weight_name(key: str) -> str: + if key.startswith("emb."): + return key.replace("emb.", "model.embed_tokens.", 1) + if key.startswith("ln_out."): + return key.replace("ln_out.", "model.norm.", 1) + if key.startswith("head."): + return key.replace("head.", "lm_head.", 1) + + parts = key.split(".") + if len(parts) < 4 or parts[0] != "blocks": + raise ValueError(f"Unsupported canonical RWKV-7 tensor name `{key}`.") + + layer_idx, component, *suffix = parts[1:] + if component == "ln0": + if layer_idx != "0": + raise ValueError(f"Unsupported canonical RWKV-7 tensor name `{key}`.") + return ".".join(("model", "embedding_norm", *suffix)) + component = { + "ln1": "input_layernorm", + "ln2": "post_attention_layernorm", + "att": "linear_attn", + "ffn": "mlp", + }.get(component) + if component is None: + raise ValueError(f"Unsupported canonical RWKV-7 tensor name `{key}`.") + if component == "linear_attn": + suffix[0] = { + "receptance": "r_proj", + "key": "k_proj", + "value": "v_proj", + "output": "o_proj", + "ln_x": "g_norm", + }.get(suffix[0], suffix[0]) + elif component == "mlp": + suffix[0] = {"key": "up_proj", "value": "down_proj"}.get(suffix[0], suffix[0]) + return ".".join(("model", "layers", layer_idx, component, *suffix)) + + def convert_state_dict(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: converted = {} for key, tensor in state.items(): - target_key = key if key.startswith("head.") else f"model.{key}" + target_key = _transformers_weight_name(key) # Released `.pth` checkpoints can store contiguous tensor views backed by # a larger flat storage. Clone so each Safetensors entry owns exactly its # logical bytes instead of retaining the source storage span. diff --git a/tests/models/rwkv/test_modeling_rwkv.py b/tests/models/rwkv/test_modeling_rwkv.py index d09a1f609b65..d2b4ba7cc438 100644 --- a/tests/models/rwkv/test_modeling_rwkv.py +++ b/tests/models/rwkv/test_modeling_rwkv.py @@ -38,11 +38,11 @@ if importlib.util.find_spec("torch") is not None: import torch - from transformers import RwkvCache, RwkvForCausalLM, RwkvModel, RwkvTimeMix, RwkvTrainingState + from transformers import RwkvCache, RwkvForCausalLM, RwkvLinearAttention, RwkvModel, RwkvTrainingState from transformers.integrations.flash_rwkv2 import _INFERENCE_OPERATORS from transformers.models.rwkv.modeling_rwkv import ( _cache_states, - _infer_tmix_projection_spec, + _infer_linear_attention_projection_spec, _load_flash_rwkv2, _stateful_training_metadata, _validate_rwkv_attention_mask, @@ -95,36 +95,72 @@ def train_temp_tmix_init(module, config, layer_idx: int) -> None: parameter = getattr(module, name) gain = math.sqrt(parameter.shape[0] / parameter.shape[1]) if parameter.shape[0] > parameter.shape[1] else 1 torch.nn.init.orthogonal_(parameter, gain=gain * 0.1) - torch.nn.init.orthogonal_(module.receptance.weight, gain=1.0) - torch.nn.init.orthogonal_(module.key.weight, gain=0.1) - torch.nn.init.orthogonal_(module.value.weight, gain=1.0) - module.output.weight.zero_() - module.ln_x.weight.fill_(((layer_idx + 1) / config.num_hidden_layers) ** 0.7) - module.ln_x.bias.zero_() + torch.nn.init.orthogonal_(module.r_proj.weight, gain=1.0) + torch.nn.init.orthogonal_(module.k_proj.weight, gain=0.1) + torch.nn.init.orthogonal_(module.v_proj.weight, gain=1.0) + module.o_proj.weight.zero_() + module.g_norm.weight.fill_(((layer_idx + 1) / config.num_hidden_layers) ** 0.7) + module.g_norm.bias.zero_() def train_temp_model_init(model) -> None: config = model.config with torch.no_grad(): - model.model.emb.weight.uniform_(-1e-4, 1e-4) - for layer_idx, block in enumerate(model.model.blocks): - for layer_norm in (getattr(block, "ln0", None), block.ln1, block.ln2): - if layer_norm is not None: - layer_norm.weight.fill_(1.0) - layer_norm.bias.zero_() - train_temp_tmix_init(block.att, config, layer_idx) + model.model.embed_tokens.weight.uniform_(-1e-4, 1e-4) + model.model.embedding_norm.weight.fill_(1.0) + model.model.embedding_norm.bias.zero_() + for layer_idx, layer in enumerate(model.model.layers): + for layer_norm in (layer.input_layernorm, layer.post_attention_layernorm): + layer_norm.weight.fill_(1.0) + layer_norm.bias.zero_() + train_temp_tmix_init(layer.linear_attn, config, layer_idx) channels = config.hidden_size ratio_1_to_almost0 = 1.0 - layer_idx / config.num_hidden_layers ddd = torch.arange(channels, dtype=torch.float32).view(1, 1, -1) / channels - block.ffn.x_k.copy_(1.0 - ddd.pow(ratio_1_to_almost0**4)) - torch.nn.init.orthogonal_(block.ffn.key.weight, gain=1.0) - block.ffn.value.weight.zero_() - model.model.ln_out.weight.fill_(1.0) - model.model.ln_out.bias.zero_() + layer.mlp.x_k.copy_(1.0 - ddd.pow(ratio_1_to_almost0**4)) + torch.nn.init.orthogonal_(layer.mlp.up_proj.weight, gain=1.0) + layer.mlp.down_proj.weight.zero_() + model.model.norm.weight.fill_(1.0) + model.model.norm.bias.zero_() gain = ( 0.5 * math.sqrt(config.vocab_size / config.hidden_size) if config.vocab_size > config.hidden_size else 0.5 ) - torch.nn.init.orthogonal_(model.head.weight, gain=gain) + torch.nn.init.orthogonal_(model.lm_head.weight, gain=gain) + + +def canonical_rwkv7_state_dict(model) -> dict[str, torch.Tensor]: + canonical = {} + for key, value in model.state_dict().items(): + if key.startswith("model.embed_tokens."): + source_key = key.replace("model.embed_tokens.", "emb.", 1) + elif key.startswith("model.embedding_norm."): + source_key = key.replace("model.embedding_norm.", "blocks.0.ln0.", 1) + elif key.startswith("model.norm."): + source_key = key.replace("model.norm.", "ln_out.", 1) + elif key.startswith("lm_head."): + source_key = key.replace("lm_head.", "head.", 1) + else: + parts = key.split(".") + layer_idx, component, *suffix = parts[2:] + component = { + "input_layernorm": "ln1", + "post_attention_layernorm": "ln2", + "linear_attn": "att", + "mlp": "ffn", + }[component] + if component == "att": + suffix[0] = { + "r_proj": "receptance", + "k_proj": "key", + "v_proj": "value", + "o_proj": "output", + "g_norm": "ln_x", + }.get(suffix[0], suffix[0]) + elif component == "ffn": + suffix[0] = {"up_proj": "key", "down_proj": "value"}.get(suffix[0], suffix[0]) + source_key = ".".join(("blocks", layer_idx, component, *suffix)) + canonical[source_key] = value + return canonical @require_torch @@ -188,26 +224,26 @@ def test_default_contract(self): def test_training_state_contract_and_selective_reset(self): config = tiny_config() state = RwkvTrainingState.zeros(config, 3, device="cpu", dtype=torch.bfloat16) - self.assertEqual(state.time_mix_shift.shape, (2, 3, 128)) - self.assertEqual(state.wkv.shape, (2, 3, 2, 64, 64)) - self.assertEqual(state.wkv.dtype, torch.float32) + self.assertEqual(state.attention_shift.shape, (2, 3, 128)) + self.assertEqual(state.recurrent_state.shape, (2, 3, 2, 64, 64)) + self.assertEqual(state.recurrent_state.dtype, torch.float32) state.validate(config, batch_size=3, device="cpu", dtype=torch.bfloat16) - state.time_mix_shift.fill_(1) - state.wkv.fill_(2) - state.channel_mix_shift.fill_(3) + state.attention_shift.fill_(1) + state.recurrent_state.fill_(2) + state.mlp_shift.fill_(3) cloned = state.clone_detach() - cloned.reset_([1], wkv=False) - self.assertEqual(cloned.time_mix_shift[:, 1].count_nonzero(), 0) - self.assertEqual(cloned.channel_mix_shift[:, 1].count_nonzero(), 0) - self.assertTrue(torch.all(cloned.wkv == 2)) - self.assertTrue(torch.all(state.time_mix_shift == 1)) + cloned.reset_([1], recurrent=False) + self.assertEqual(cloned.attention_shift[:, 1].count_nonzero(), 0) + self.assertEqual(cloned.mlp_shift[:, 1].count_nonzero(), 0) + self.assertTrue(torch.all(cloned.recurrent_state == 2)) + self.assertTrue(torch.all(state.attention_shift == 1)) def test_training_state_rejects_shape_dtype_and_device_mismatch(self): config = tiny_config() state = RwkvTrainingState.zeros(config, 2, device="cpu", dtype=torch.bfloat16) with self.assertRaisesRegex(ValueError, "must have shape"): state.clone().validate(config, batch_size=1, device="cpu", dtype=torch.bfloat16) - state.wkv = state.wkv.to(torch.bfloat16) + state.recurrent_state = state.recurrent_state.to(torch.bfloat16) with self.assertRaisesRegex(TypeError, "torch.float32"): state.validate(config, batch_size=2, device="cpu", dtype=torch.bfloat16) @@ -284,17 +320,23 @@ def test_auto_mappings_keep_rwkv_identity(self): class Rwkv7ModelStructureTest(unittest.TestCase): all_model_classes = (RwkvModel, RwkvForCausalLM) - def test_public_component_and_canonical_tensor_names(self): + def test_public_component_and_transformers_tensor_names(self): model = RwkvForCausalLM(tiny_config()) - self.assertIsInstance(model.model.blocks[0].att, RwkvTimeMix) + self.assertIsInstance(model.model.layers[0].linear_attn, RwkvLinearAttention) state = model.state_dict() - self.assertEqual(state["model.blocks.0.att.w1"].shape, (128, 32)) - self.assertEqual(state["model.blocks.0.att.w2"].shape, (32, 128)) - self.assertEqual(state["model.blocks.0.att.v0"].shape, (1, 1, 128)) - self.assertEqual(state["model.blocks.0.att.v1"].shape, (128, 32)) - self.assertEqual(state["model.blocks.1.att.v1"].shape, (128, 32)) - self.assertEqual(state["model.blocks.0.ffn.key.weight"].shape, (512, 128)) - self.assertEqual(state["model.blocks.0.ffn.value.weight"].shape, (128, 512)) + self.assertEqual(state["model.layers.0.linear_attn.w1"].shape, (128, 32)) + self.assertEqual(state["model.layers.0.linear_attn.w2"].shape, (32, 128)) + self.assertEqual(state["model.layers.0.linear_attn.v0"].shape, (1, 1, 128)) + self.assertEqual(state["model.layers.0.linear_attn.v1"].shape, (128, 32)) + self.assertEqual(state["model.layers.1.linear_attn.v1"].shape, (128, 32)) + self.assertEqual(state["model.layers.0.linear_attn.r_proj.weight"].shape, (128, 128)) + self.assertEqual(state["model.layers.0.linear_attn.g_norm.weight"].shape, (128,)) + self.assertEqual(state["model.layers.0.mlp.up_proj.weight"].shape, (512, 128)) + self.assertEqual(state["model.layers.0.mlp.down_proj.weight"].shape, (128, 512)) + self.assertEqual(state["model.embed_tokens.weight"].shape, (256, 128)) + self.assertEqual(state["model.embedding_norm.weight"].shape, (128,)) + self.assertEqual(state["model.norm.weight"].shape, (128,)) + self.assertEqual(state["lm_head.weight"].shape, (256, 128)) self.assertTrue(all(tensor.isfinite().all() for tensor in state.values())) def test_runtime_fails_closed_on_cpu(self): @@ -325,19 +367,19 @@ def test_model_rejects_padding_before_loading_the_provider(self): use_cache=False, ) - def test_time_mix_and_channel_mix_share_attention_mask_validation(self): - block = RwkvForCausalLM(tiny_config()).model.blocks[0] - hidden_states = torch.zeros(1, 2, block.att.config.hidden_size) + def test_linear_attention_and_mlp_share_attention_mask_validation(self): + layer = RwkvForCausalLM(tiny_config()).model.layers[0] + hidden_states = torch.zeros(1, 2, layer.linear_attn.config.hidden_size) invalid_mask = torch.ones(1, 2, 1) - for layer in (block.att, block.ffn): - with self.subTest(layer=type(layer).__name__), self.assertRaisesRegex(ValueError, "two-dimensional"): - layer(hidden_states, attention_mask=invalid_mask) + for module in (layer.linear_attn, layer.mlp): + with self.subTest(module=type(module).__name__), self.assertRaisesRegex(ValueError, "two-dimensional"): + module(hidden_states, attention_mask=invalid_mask) def test_explicit_low_rank_dimensions_control_parameter_shapes(self): model = RwkvForCausalLM( tiny_config(decay_low_rank_dim=17, a_low_rank_dim=19, v_low_rank_dim=23, gate_low_rank_dim=29) ) - attention = model.model.blocks[0].att + attention = model.model.layers[0].linear_attn self.assertEqual(attention.w1.shape, (128, 17)) self.assertEqual(attention.a1.shape, (128, 19)) self.assertEqual(attention.v1.shape, (128, 23)) @@ -349,7 +391,7 @@ def test_initialization_matches_train_temp_tensor_by_tensor(self): expected = RwkvForCausalLM(config) torch.manual_seed(20260806) actual.model.reset_parameters() - actual.reset_head_parameters() + actual.reset_lm_head_parameters() torch.manual_seed(20260806) train_temp_model_init(expected) for name, tensor in actual.state_dict().items(): @@ -447,9 +489,9 @@ def test_declares_standard_gradient_checkpointing_support(self): self.assertTrue(model.supports_gradient_checkpointing) self.assertFalse(model.supports_tp_plan) self.assertFalse(model._can_compile_fullgraph) - self.assertTrue(all(block.gradient_checkpointing for block in model.model.blocks)) + self.assertTrue(all(layer.gradient_checkpointing for layer in model.model.layers)) - def test_cache_state_shape_is_owned_by_each_time_mix_layer(self): + def test_cache_state_shape_is_owned_by_each_linear_attention_layer(self): cache_config = mock.Mock(num_hidden_layers=2, number_of_conv_states=2) cache = RwkvCache(cache_config) hidden_states = torch.zeros(1, 3, 512) @@ -544,10 +586,7 @@ def test_build_tokenizer_requires_rwkv_trie_from_file(self): def test_minimal_conversion_and_fresh_model_process(self): converter = self._converter_module() model = RwkvForCausalLM(tiny_config()) - source = { - key.removeprefix("model.") if key.startswith("model.") else key: value - for key, value in model.state_dict().items() - } + source = canonical_rwkv7_state_dict(model) with tempfile.TemporaryDirectory() as directory: root = Path(directory) checkpoint = root / "rwkv7.pth" @@ -583,10 +622,7 @@ def test_converter_preserves_nonformula_low_rank_dimensions(self): converter = self._converter_module() config = tiny_config(decay_low_rank_dim=17, a_low_rank_dim=19, v_low_rank_dim=23, gate_low_rank_dim=29) model = RwkvForCausalLM(config) - source = { - key.removeprefix("model.") if key.startswith("model.") else key: value - for key, value in model.state_dict().items() - } + source = canonical_rwkv7_state_dict(model) inferred = converter.infer_config(source, context_length=16) self.assertEqual( ( @@ -601,10 +637,7 @@ def test_converter_preserves_nonformula_low_rank_dimensions(self): def test_converter_rejects_cross_layer_low_rank_drift(self): converter = self._converter_module() model = RwkvForCausalLM(tiny_config()) - source = { - key.removeprefix("model.") if key.startswith("model.") else key: value - for key, value in model.state_dict().items() - } + source = canonical_rwkv7_state_dict(model) source["blocks.1.att.w1"] = torch.empty(128, 31) source["blocks.1.att.w2"] = torch.empty(31, 128) with self.assertRaisesRegex(ValueError, "must agree across all layers"): @@ -644,7 +677,7 @@ def test_converter_rejects_mixed_orig_mod_prefix(self): def test_converter_detaches_larger_source_storage(self): source_storage = torch.arange(16, dtype=torch.float32) converted = self._converter_module().convert_state_dict({"emb.weight": source_storage[4:8]}) - tensor = converted["model.emb.weight"] + tensor = converted["model.embed_tokens.weight"] self.assertEqual(tensor.untyped_storage().nbytes(), tensor.numel() * tensor.element_size()) self.assertEqual(tensor.tolist(), [4.0, 5.0, 6.0, 7.0]) @@ -659,7 +692,7 @@ def test_unmerged_lora_projection_spec_matches_peft(self): config = tiny_config(hidden_size=1024, intermediate_size=4096) model = RwkvForCausalLM(config).cuda().eval() - targets = ["receptance", "key", "value", "output"] + targets = ["r_proj", "k_proj", "v_proj", "o_proj"] model.add_adapter( LoraConfig(r=8, lora_alpha=16, target_modules=targets, init_lora_weights=False), adapter_name="first", @@ -669,10 +702,10 @@ def test_unmerged_lora_projection_spec_matches_peft(self): x = torch.randn(5, config.hidden_size, device="cuda", dtype=torch.float16) for name in targets: - projection = getattr(model.model.blocks[0].att, name) + projection = getattr(model.model.layers[0].linear_attn, name) with torch.no_grad(): expected = projection(x) - weight, lora_a, lora_b, scale = _infer_tmix_projection_spec(projection) + weight, lora_a, lora_b, scale = _infer_linear_attention_projection_spec(projection) actual = torch.nn.functional.linear(x, weight) if lora_a is not None: actual = actual + torch.nn.functional.linear(torch.nn.functional.linear(x, lora_a), lora_b) * scale @@ -683,10 +716,10 @@ def test_unmerged_lora_projection_spec_matches_peft(self): output = model(input_ids, use_cache=False) self.assertTrue(torch.isfinite(output.logits).all()) - projection = model.model.blocks[0].att.receptance + projection = model.model.layers[0].linear_attn.r_proj model.disable_adapters() with torch.no_grad(): - weight, lora_a, lora_b, scale = _infer_tmix_projection_spec(projection) + weight, lora_a, lora_b, scale = _infer_linear_attention_projection_spec(projection) disabled = torch.nn.functional.linear(x, weight) base = projection.get_base_layer()(x) self.assertIsNone(lora_a) @@ -701,7 +734,7 @@ def test_unmerged_lora_projection_spec_matches_peft(self): ) model.set_adapter(["first", "second"]) with self.assertRaisesRegex(RuntimeError, "exactly one active"): - _infer_tmix_projection_spec(projection) + _infer_linear_attention_projection_spec(projection) def test_stateful_chunk_forward_matches_single_stateful_call(self): config = tiny_config() @@ -727,23 +760,23 @@ def test_stateful_training_backward_and_fp32_wkv_contract(self): input_ids = torch.randint(0, config.vocab_size, (1, 7), device="cuda") state = RwkvTrainingState.zeros(config, 1, device="cuda", dtype=torch.bfloat16) outputs = model(input_ids, labels=input_ids, training_state=state, use_cache=False) - self.assertEqual(outputs.training_state.wkv.dtype, torch.float32) + self.assertEqual(outputs.training_state.recurrent_state.dtype, torch.float32) outputs.loss.backward() - self.assertIsNotNone(model.model.blocks[0].att.receptance.weight.grad) - self.assertTrue(torch.isfinite(model.model.blocks[0].att.receptance.weight.grad).all()) + self.assertIsNotNone(model.model.layers[0].linear_attn.r_proj.weight.grad) + self.assertTrue(torch.isfinite(model.model.layers[0].linear_attn.r_proj.weight.grad).all()) def test_gradient_checkpointing_matches_stateless_and_stateful_training(self): config = tiny_config() model = RwkvForCausalLM(config).cuda().to(torch.bfloat16).train() with torch.no_grad(): - for block in model.model.blocks: - block.att.output.weight.normal_(std=0.01) - block.ffn.value.weight.normal_(std=0.01) + for layer in model.model.layers: + layer.linear_attn.o_proj.weight.normal_(std=0.01) + layer.mlp.down_proj.weight.normal_(std=0.01) input_ids = torch.randint(0, config.vocab_size, (1, 16), device="cuda") gradient_names = ( - "model.blocks.0.att.receptance.weight", - "model.blocks.0.att.output.weight", - "head.weight", + "model.layers.0.linear_attn.r_proj.weight", + "model.layers.0.linear_attn.o_proj.weight", + "lm_head.weight", ) def run(*, checkpointing: bool, stateful: bool): @@ -787,44 +820,44 @@ def run(*, checkpointing: bool, stateful: bool): for name in gradient_names: torch.testing.assert_close(actual_gradients[name], expected_gradients[name], atol=0, rtol=0) - def test_inference_preparation_offloads_only_canonical_ffn_down_layout(self): + def test_inference_preparation_offloads_only_mlp_down_projection(self): model = RwkvForCausalLM(tiny_config(hidden_size=1024, intermediate_size=4096)).cuda().eval() - expected = model.model.blocks[0].ffn.value.weight.detach().cpu().half().clone() + expected = model.model.layers[0].mlp.down_proj.weight.detach().cpu().half().clone() model.prepare_for_inference().prepare_for_inference() - channel_mix = model.model.blocks[0].ffn - self.assertEqual(channel_mix.value.weight.device.type, "cpu") - self.assertEqual(channel_mix._value_runtime.device.type, "cuda") - torch.testing.assert_close(channel_mix.value.weight, expected, atol=0, rtol=0) + mlp = model.model.layers[0].mlp + self.assertEqual(mlp.down_proj.weight.device.type, "cpu") + self.assertEqual(mlp._down_proj_runtime.device.type, "cuda") + torch.testing.assert_close(mlp.down_proj.weight, expected, atol=0, rtol=0) with torch.no_grad(): - channel_mix.value.weight.add_(1) + mlp.down_proj.weight.add_(1) model.prepare_for_inference() torch.testing.assert_close( - channel_mix._value_runtime, - channel_mix.value.weight.T.cuda(), + mlp._down_proj_runtime, + mlp.down_proj.weight.T.cuda(), atol=0, rtol=0, ) - expected = channel_mix.value.weight.detach().clone() + expected = mlp.down_proj.weight.detach().clone() with tempfile.TemporaryDirectory() as directory: model.save_pretrained(directory) reloaded = RwkvForCausalLM.from_pretrained(directory, dtype=torch.float16) - torch.testing.assert_close(reloaded.model.blocks[0].ffn.value.weight, expected, atol=0, rtol=0) - self.assertIsNone(reloaded.model.blocks[0].ffn._value_runtime) + torch.testing.assert_close(reloaded.model.layers[0].mlp.down_proj.weight, expected, atol=0, rtol=0) + self.assertIsNone(reloaded.model.layers[0].mlp._down_proj_runtime) model.train() input_ids = torch.randint(0, model.config.vocab_size, (1, 16), device="cuda") with self.assertRaisesRegex(RuntimeError, "in-place Albatross inference layout"): model(input_ids, labels=input_ids, use_cache=False) model.to(device="cuda", dtype=torch.bfloat16).train() - self.assertTrue(all(block.ffn._value_runtime is None for block in model.model.blocks)) - self.assertTrue(all(block.att._w1_original is None for block in model.model.blocks)) - self.assertTrue(all(block.ffn.value.weight.is_cuda for block in model.model.blocks)) + self.assertTrue(all(layer.mlp._down_proj_runtime is None for layer in model.model.layers)) + self.assertTrue(all(layer.linear_attn._w1_original is None for layer in model.model.layers)) + self.assertTrue(all(layer.mlp.down_proj.weight.is_cuda for layer in model.model.layers)) outputs = model(input_ids, labels=input_ids, use_cache=False) outputs.loss.backward() - self.assertIsNotNone(model.model.blocks[0].ffn.value.weight.grad) - self.assertTrue(torch.isfinite(model.model.blocks[0].ffn.value.weight.grad).all()) + self.assertIsNotNone(model.model.layers[0].mlp.down_proj.weight.grad) + self.assertTrue(torch.isfinite(model.model.layers[0].mlp.down_proj.weight.grad).all()) def test_recurrent_metadata_is_scoped_to_cuda_stream(self): class FakeFlashRwkv2: @@ -881,8 +914,8 @@ def test_training_forward_backward_uses_public_train_temp_family(self): self.assertEqual(outputs.logits.dtype, torch.bfloat16) self.assertTrue(torch.isfinite(outputs.loss)) outputs.loss.backward() - self.assertIsNotNone(model.model.blocks[0].att.output.weight.grad) - self.assertGreater(model.model.blocks[0].att.output.weight.grad.abs().max().item(), 0) + self.assertIsNotNone(model.model.layers[0].linear_attn.o_proj.weight.grad) + self.assertGreater(model.model.layers[0].linear_attn.o_proj.weight.grad.abs().max().item(), 0) def test_outputs_use_one_rwkv_type_and_standard_loss_hook(self): config = tiny_config() @@ -930,7 +963,7 @@ def test_inference_prefill_decode_and_continuation(self): staged.past_key_values.batch_select_indices(torch.tensor([1], device="cuda")) self.assertEqual(staged.past_key_values.batch_size, 1) - def test_inference_output_hidden_states_preserves_block_boundaries(self): + def test_inference_output_hidden_states_preserves_layer_boundaries(self): config = tiny_config(hidden_size=1024, intermediate_size=4096) model = RwkvForCausalLM(config).cuda().eval().prepare_for_inference() input_ids = torch.randint(0, config.vocab_size, (1, 5), device="cuda") diff --git a/utils/not_doctested.txt b/utils/not_doctested.txt index f3aa6511c058..4b0c2fbbfdc3 100644 --- a/utils/not_doctested.txt +++ b/utils/not_doctested.txt @@ -575,7 +575,6 @@ src/transformers/models/roberta_prelayernorm/convert_roberta_prelayernorm_origin src/transformers/models/roc_bert/configuration_roc_bert.py src/transformers/models/roformer/modeling_roformer.py src/transformers/models/rwkv/configuration_rwkv.py -src/transformers/models/rwkv/convert_rwkv_checkpoint_to_hf.py src/transformers/models/rwkv/modeling_rwkv.py src/transformers/models/sam/configuration_sam.py src/transformers/models/sam/convert_sam_to_hf.py