Usage
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
BASE = "Jackrong/Qwen3.5-35B-A3B-Claude-4.6-Opus-Reasoning-Distilled"
ADAPTER = "thanhdath/Qwen3.5-35B-A3B-Legal-LoRA"
tok = AutoTokenizer.from_pretrained(ADAPTER)
model = AutoModelForCausalLM.from_pretrained(BASE, torch_dtype=torch.bfloat16, device_map="auto")
model = PeftModel.from_pretrained(model, ADAPTER)
msgs = [{"role": "user", "content": "Tóm tắt các điểm chính của hợp đồng lao động này ..."}]
inputs = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt").to(model.device)
out = model.generate(inputs, max_new_tokens=1024)
print(tok.decode(out[0][inputs.shape[1]:], skip_special_tokens=True))
To produce a standalone (merged) full model, load the base in bf16 (not quantized), attach this
adapter, and call model.merge_and_unload() before save_pretrained(...).
Training configuration
Table | |
|---|
| Method | LoRA (PEFT) SFT, bf16, no quantization / no offload |
| Base model | Jackrong/Qwen3.5-35B-A3B-Claude-4.6-Opus-Reasoning-Distilled (Qwen3.5 MoE, 256 experts, 3B active) |
| Hardware | 1× NVIDIA A100 80 GB — 69 GB base fits fully on-GPU |
| Epochs | 1 (266 optimizer steps over 8,500 samples) |
| Global batch | 32 (per-device 1 × grad-accum 32) |
| Max sequence length | 8192 (covers the assistant span of ~89% of samples) |
| Optimizer / LR | AdamW, 2e-5, cosine decay, warmup ratio 0.03 |
| Loss | assistant-only ({% generation %} masked) |
| LoRA | r=16, α=32, dropout=0.05 |
| LoRA targets | attention q/k/v/o_proj, linear-attn in_proj_qkv/a/b/z + out_proj, shared-expert gate/up/down_proj |
| Attention | flash-attention-2 (full-attn) + flash-linear-attention (linear-attn) |
| MoE aux loss | disabled (router_aux_loss_coef=0; the router is frozen under LoRA) |
| Gradient checkpointing | on (non-reentrant) |
| Wall-clock | ~7.5 h |
Note on LoRA scope: the 256 routed experts are stored as fused Qwen3_5MoeExperts tensors
(~93% of parameters) and are not nn.Linear modules, so LoRA targets the attention projections and
the shared expert MLP (which every token passes through). Liger kernels are disabled because the
fused-MoE backward kernel does not compile on Triton 3.2; TRL's chunked cross-entropy handles the
248k-token vocabulary head instead.
Loss curve (train)
Table with columns: step, loss, step, loss| step | loss | | step | loss |
|---|
| 1 | 0.409 | | 160 | 0.306 |
| 40 | 0.329 | | 200 | 0.283 |
| 80 | 0.306 | | 240 | 0.319 |
| 120 |
Full per-step history (loss / grad-norm / learning-rate), the raw console log, the exact PBS job
script and the training chat-template source are in logs/.
Dataset
Curated "daily-task" mix (~8.5k conversations, seed 42) from the CATI-AI legal SFT collection:
closed QA, markdown QA pairs, legal-document summarization, legal-conflict distillation, judicial
reasoning, objective legal-opinion generation, document generation, and user-intent understanding.
Pure-benchmark sources were dropped.
Framework versions
- PEFT 0.19.1
- TRL 1.7.1
- Transformers 5.13.0
- PyTorch 2.6.0+cu124
- Datasets 4.8.5
- Tokenizers 0.22.2
Citation
@software{vonwerra2020trl,
title = {{TRL: Transformers Reinforcement Learning}},
author = {von Werra, Leandro and Belkada, Younes and Tunstall, Lewis and Beeching, Edward and Thrush, Tristan and Lambert, Nathan and Huang, Shengyi and Rasul, Kashif and Gallouédec, Quentin},
license = {Apache-2.0},
url = {https://github.com/huggingface/trl}
}