Add TRL SFT compatibility

README.md

@@ -213,6 +213,44 @@ Reference stops are token ID `0` and `\n\nUser:`.
 Strip trailing spaces from user input. The official RWKV prompt guide is available
 in [`RWKV7-G1x-templates.txt`](https://github.com/BlinkDL/RWKV-LM/blob/main/RWKV-v7/RWKV7-G1x-templates.txt).
 
+## Supervised fine-tuning
+
+The bundled model supports TRL 1.10+ `SFTTrainer`, including its default
+`chunked_nll`, gradient checkpointing, assistant-only loss, BFD packing, and PEFT
+LoRA. Packing boundaries carried as reset `position_ids` are converted into RWKV
+recurrent-state boundaries. Do not use the boundary-destroying `wrapped` packing
+strategy.
+
+```python
+from datasets import load_dataset
+from peft import LoraConfig
+from trl import SFTConfig, SFTTrainer
+
+# Reuse `model` and `tokenizer` loaded in the Transformers quickstart above.
+dataset = load_dataset("trl-lib/Capybara", split="train")
+trainer = SFTTrainer(
+    model=model,
+    processing_class=tokenizer,
+    train_dataset=dataset,
+    args=SFTConfig(
+        output_dir="rwkv7-sft",
+        max_length=2048,
+        packing=True,
+        packing_strategy="bfd",
+        assistant_only_loss=True,
+        use_cache=False,
+        gradient_checkpointing=True,
+    ),
+    peft_config=LoraConfig(
+        task_type="CAUSAL_LM",
+        r=8,
+        lora_alpha=16,
+        target_modules=["receptance", "key", "value", "output"],
+    ),
+)
+trainer.train()
+```
+
 ## Optimized local inference
 
 Install the versions listed in `inference/requirements.txt`, then run the bundled

chat_template.jinja

@@ -1,4 +1,4 @@
-{%- set add_generation_prompt = add_generation_prompt | default(true) -%}
+{%- set add_generation_prompt = add_generation_prompt | default(false) -%}
 {%- set thinking = thinking | default(false) -%}
 {%- set bos_token = bos_token | default('', true) -%}
 {%- set tools = tools | default([], true) -%}
@@ -23,6 +23,7 @@
 {%- if message.role == 'user' -%}
 {{ 'User: ' ~ (message.content | trim) ~ '\n\n' }}
 {%- elif message.role == 'assistant' -%}
+{%- generation -%}
 {%- set content = message.content | default('', true) | trim -%}
 {{ 'Assistant:' }}
 {%- if message.tool_calls is defined and message.tool_calls | length > 0 -%}
@@ -43,6 +44,7 @@
 {{ ' ' ~ content }}
 {%- endif -%}
 {{ '\n\n' }}
+{%- endgeneration -%}
 {%- elif message.role == 'tool' -%}
 {{ 'User: Function output:\n' ~ (message.content | trim) ~ '\n\n' }}
 {%- endif -%}

modeling_rwkv7.py

@@ -407,9 +407,14 @@ class Rwkv7Attention(nn.Module):
         # The value-residual LoRA exists on every layer in a reference checkpoint,
         # but layer 0 PRODUCES v_first instead of mixing towards it, so its copy is
         # never read. It is registered anyway so that loading is lossless.
-        self.v1 = nn.Parameter(torch.zeros(C, config.v_low_rank_dim))
-        self.v2 = nn.Parameter(torch.zeros(config.v_low_rank_dim, C))
-        self.v0 = nn.Parameter(torch.zeros(1, 1, C))
+        if layer_id == 0:
+            self.register_buffer("v1", torch.zeros(C, config.v_low_rank_dim))
+            self.register_buffer("v2", torch.zeros(config.v_low_rank_dim, C))
+            self.register_buffer("v0", torch.zeros(1, 1, C))
+        else:
+            self.v1 = nn.Parameter(torch.zeros(C, config.v_low_rank_dim))
+            self.v2 = nn.Parameter(torch.zeros(config.v_low_rank_dim, C))
+            self.v0 = nn.Parameter(torch.zeros(1, 1, C))
 
         self.k_k = nn.Parameter(torch.zeros(1, 1, C))
         self.k_a = nn.Parameter(torch.zeros(1, 1, C))
@@ -814,6 +819,7 @@ class Rwkv7Output(ModelOutput):
 
     last_hidden_state: torch.FloatTensor | None = None
     state: Rwkv7Cache | None = None
+    past_key_values: None = None
     hidden_states: tuple[torch.FloatTensor, ...] | None = None
     attentions: None = None
 
@@ -906,6 +912,19 @@ class Rwkv7PreTrainedModel(PreTrainedModel):
             init.uniform_(module.weight, a=-1e-4, b=1e-4)
 
 
+def _packed_boundaries_from_position_ids(
+    position_ids: torch.LongTensor | None, seq_len: int
+) -> torch.LongTensor | None:
+    if position_ids is None or position_ids.ndim != 2 or position_ids.shape[0] != 1:
+        return None
+    starts = torch.nonzero(position_ids[0] == 0, as_tuple=False).flatten()
+    if starts.numel() <= 1:
+        return None
+    if starts[0].item() != 0:
+        raise ValueError("packed position_ids must start at zero")
+    return torch.cat((starts, starts.new_tensor([seq_len])))
+
+
 @auto_docstring
 class Rwkv7Model(Rwkv7PreTrainedModel):
     def __init__(self, config: Rwkv7Config):
@@ -948,6 +967,7 @@ class Rwkv7Model(Rwkv7PreTrainedModel):
         self,
         input_ids: torch.LongTensor | None = None,
         attention_mask: torch.LongTensor | None = None,
+        position_ids: torch.LongTensor | None = None,
         inputs_embeds: torch.FloatTensor | None = None,
         state: Rwkv7Cache | None = None,
         cu_seq_lens: torch.LongTensor | None = None,
@@ -1005,6 +1025,11 @@ class Rwkv7Model(Rwkv7PreTrainedModel):
         if inputs_embeds is None:
             inputs_embeds = self.emb(input_ids)
 
+        if cu_seq_lens is None and attention_mask is None:
+            cu_seq_lens = _packed_boundaries_from_position_ids(
+                position_ids, inputs_embeds.shape[1]
+            )
+
         if use_cache and state is None:
             state = self._empty_state(inputs_embeds.shape[0], inputs_embeds.device, inputs_embeds.dtype)
 
@@ -1137,6 +1162,7 @@ class Rwkv7ForCausalLM(Rwkv7PreTrainedModel, GenerationMixin):
         self,
         input_ids: torch.LongTensor | None = None,
         attention_mask: torch.LongTensor | None = None,
+        position_ids: torch.LongTensor | None = None,
         inputs_embeds: torch.FloatTensor | None = None,
         state: Rwkv7Cache | None = None,
         labels: torch.LongTensor | None = None,
@@ -1159,9 +1185,12 @@ class Rwkv7ForCausalLM(Rwkv7PreTrainedModel, GenerationMixin):
             so `generate` declined to pass it and a caller who passed it was quietly
             ignored.
         """
+        if labels is not None and use_cache is None:
+            use_cache = False
         outputs = self.rwkv7(
             input_ids=input_ids,
             attention_mask=attention_mask,
+            position_ids=position_ids,
             inputs_embeds=inputs_embeds,
             state=state,
             use_cache=use_cache,

release-manifest.json

@@ -1,7 +1,7 @@
 {
   "builder": {
     "asset_set": "2026.08.06-r1",
-    "assets_sha256": "5a31129dcca03f5ecd55157e1561af64aaf35046c77e253fc4d6214b95ab8ede",
+    "assets_sha256": "536d0d0afa96630b2bad11534ebe35434226f83b8999a1c60d76fe37802ff139",
     "version": "0.2.0"
   },
   "conversion": {
@@ -36,13 +36,13 @@
     },
     "README.md": {
       "role": "model_card",
-      "sha256": "0f1a8bfe6a0d420333fdf5f7e787e751104348c81c531e3b846a6b0bf7b88d44",
-      "size_bytes": 10331
+      "sha256": "fbcb25fd1bb80a23a961c9bc85ff57f606177341e7d24143a5cbcc75b69413da",
+      "size_bytes": 11454
     },
     "chat_template.jinja": {
       "role": "tokenizer",
-      "sha256": "abcd52a871eb3f7f0de7b69f2e0ac974fd7cfb03d4acea7a4609cb441f5d9048",
-      "size_bytes": 2336
+      "sha256": "e0f8e7809caa4ddf6f70f89a96b76974b5d554fd3122f2f0e2023c256c71994e",
+      "size_bytes": 2378
     },
     "config.json": {
       "role": "model_config",
@@ -91,8 +91,8 @@
     },
     "modeling_rwkv7.py": {
       "role": "model_code",
-      "sha256": "f3a1d9bde755f54887451e8755c6d48866ecc9146610c7d8fa73e0fa68c6734c",
-      "size_bytes": 57446
+      "sha256": "8b8a3459b40a33424b592abff360485577a8f19419be411e3dceb817a51ebe46",
+      "size_bytes": 58685
     },
     "tokenizer.json": {
       "role": "tokenizer",
@@ -139,7 +139,13 @@
     "provenance": "locked-profile"
   },
   "model_code": {
-    "format_version": 1,
+    "format_version": 2,
+    "patches": [
+      "layer-zero-value-residual-buffers",
+      "trainer-past-key-values-placeholder",
+      "trl-position-ids-packing-boundaries",
+      "labels-disable-recurrent-cache"
+    ],
     "source_repository": "https://github.com/huggingface/transformers.git",
     "source_revision": "4ad9ed0747ed6ba75c787e8f9040dcd64b166ee2",
     "sources": {
@@ -151,7 +157,7 @@
       },
       "modeling_rwkv7.py": {
         "asset_path": "model_code/modeling_rwkv7.py",
-        "output_sha256": "f3a1d9bde755f54887451e8755c6d48866ecc9146610c7d8fa73e0fa68c6734c",
+        "output_sha256": "8b8a3459b40a33424b592abff360485577a8f19419be411e3dceb817a51ebe46",
         "repository_path": "src/transformers/models/rwkv7/modeling_rwkv7.py",
         "source_sha256": "3e8e5af7c4eba0b5de1496aef44773d7ac1bb4d96756e6f55efaf29453d67952"
       }
@@ -320,7 +326,7 @@
       }
     }
   },
-  "schema_version": 7,
+  "schema_version": 8,
   "source": {
     "filename": "rwkv7-g1i-1.5b-20260805-ctx16384.pth",
     "kind": "huggingface",