Jadson/Nanbeige4.2-3B-dlora

Nanbeige4.2-3B dlora - diffuser checkpoint (training repo)

This repo carries the trained diffuser of the dlora project, ready for continued training (SFT of chat). The diffuser is a small denoiser (22.0M parameters) that, for every token, generates a rank-16 LoRA delta for the FFN weights of a fully frozen backbone (Nanbeige/Nanbeige4.2-3B, the instruct checkpoint, not -Base). Its input per token is only what attention produced at that step (x_attn), and it runs on all 44 FFN applications of the looped backbone (22 layers x 2 passes), sampling 4 DDIM steps per application. Only the diffuser trains; the backbone is byte-identical to the hub original.

• Inference artifact (merged GGUF, IQ4_NL with the diffuser in F32): Jadson/Nanbeige4.2-3B-dlora-GGUF
• Runtime (llama.cpp port that serves it): the fork

Files

File	What it is
difusor.ckpt	The trained diffuser: 21 fp32 tensors, 22,032,736 parameters, step 152 of the pretraining run (84.1 MiB, sha256 0048524a90679b88...). cfg inside records every hyperparameter below.
nanbeige_ffn_diffusion_lora.py	The adapter module: LoRADiffuser (DDIM), anexar_ffn_diffusion_lora (attach to a loaded backbone), restaurar_rope (mandatory after from_pretrained), ruido_da_inferencia (the fixed seed-42 inference draw).
treinar_nanbeige_ffn_diffusion_lora.py	The SFT trainer (HF/transformers + Lightning Fabric, tyro CLI): chat data via apply_chat_template, prompt tokens masked with -100, checkpointing, DDP. This is the reference recipe.

Continue training (SFT)

Dependencies: torch, transformers>=5.18, lightning, tyro, datasets, and bitsandbytes+accelerate for the NF4 backbone. The backbone is downloaded from its own hub repo; this repo only carries the trained delta.

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
from nanbeige_ffn_diffusion_lora import anexar_ffn_diffusion_lora, restaurar_rope

model = AutoModelForCausalLM.from_pretrained(
    "Nanbeige/Nanbeige4.2-3B", trust_remote_code=True, dtype=torch.bfloat16,
    quantization_config=BitsAndBytesConfig(   # the diffuser was trained against THIS backbone
        load_in_4bit=True, bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16))
model.eval()
restaurar_rope(model)   # transformers 5 leaves the remote code's RoPE inv_freq uninitialized

lora_diffuser = anexar_ffn_diffusion_lora(
    model, lora_rank=16, lora_alpha=16, latent_dim=32,
    n_diff_steps=4, entrada="x_attn")

estado = torch.load("difusor.ckpt", map_location="cpu", weights_only=False)
lora_diffuser.load_state_dict(estado["lora_diffuser"])   # fresh optimizer for the new objective

Or straight to a chat SFT run (the trainer above does the whole snippet):

python treinar_nanbeige_ffn_diffusion_lora.py \
    --model.model-id Nanbeige/Nanbeige4.2-3B --model.load-in-4bit \
    --data.dataset-name HuggingFaceH4/ultrafeedback_binarized \
    --data.train-split train_sft --data.messages-column chosen

The trainer's --resume restores the pretraining optimizer and scheduler too; for a new SFT objective initialize the weights only (the load_state_dict line), which is why difusor.ckpt carries no optimizer state.

Regime invariants

• Backbone in NF4 (--model.load-in-4bit): the diffuser was trained against the bitsandbytes NF4 backbone; loading the original bf16 weights changes the backbone under the delta and the model degenerates (measured in the project).
• Input x_attn (recorded in the checkpoint cfg).
• Frozen is not no_grad: gradients flow through the frozen backbone to the diffuser (enable_input_require_grads under checkpointing).
• restaurar_rope after every from_pretrained (transformers 5, remote code).
• Inference noise: one fixed draw, seed 42 (ruido_da_inferencia); training drew best-of-4 per document.
• Hyperparameters must match the cfg inside difusor.ckpt:

Hyperparameter	Value
backbone	Nanbeige/Nanbeige4.2-3B (frozen, NF4)
LoRA rank / alpha	16 / 16
latent dim	32
DDIM steps per FFN application	4
noise draws per document (training)	4
noise sigma	1.0

Pretraining provenance

Raw-text pretraining (no SFT): FineWeb2-HQ, 20 languages without English, whole documents packed into 262,144-token sequences, best-of-4 noise draws per document; this checkpoint is step 152 of that run.

Evaluation

Measured with the official MMLU-Pro CoT harness (5-shot, temperature 0), served through llama.cpp from the merged bf16 GGUF; the full per-category tables live in the inference repo's card. This checkpoint has not been re-evaluated after SFT - the numbers are the raw-text pretrained diffuser.

Measure	Value
Official accuracy	77.98%
Strict accuracy	77.86% (s.e. 0.38 pp)
Questions	12032
No-answer / truncated	161 / 178
Mean completion tokens	5384

Limitations

• Trained on raw text over an instruct backbone: the active delta pulls the instruct model's outputs toward pretraining statistics (the format/letter-mass effect). SFT is the intended next step, and this repo exists for it.
• No same-protocol run of the base model yet, so every number above is the dlora alone.
• The DDIM sampler adds several matmuls per FFN application (the llama.cpp graph grows ~3.5x).

License

Apache-2.0, the same as the backbone Nanbeige/Nanbeige4.2-3B.