Jadson/Nanbeige4.2-3B-dlora-sft

Nanbeige4.2-3B dlora - SFT diffuser (tool-call traces)

This repo carries the SFT-trained diffuser of the dlora project. 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, NF4). 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.

This checkpoint is step 237 of the tool-call SFT (238 steps, lr 0.0002, warmup 10, micro-batch 1 x accum 4), trained on tmp/traces.jsonl (train split, chat-format raw) — 951 raw-rendered agent traces in the Qwen template (<think>... then <tool_call><function=ipython><parameter=code>...). Loss fell 14.86 → 6.38 over the run; |decode_A| grew 0 → 11.38 monotonically (the adapter path active and healthy), and the frozen backbone was verified unchanged.

Verificacao por inferencia (amostragem temp 0.6, top_p 0.9, prompt = trace 1 ate <|im_start|>assistant + think tipico do corpus como semente de estilo, tmp/infer_sft.log):

</think>: True | <function=ipython>: True | 7393 chars
RESULT: PASS

O modelo fecha o </think>, emite <tool_call><function=ipython><parameter=code>... completo e encerra em <|im_end|>. Baseline (MODO=base, backbone puro, mesmo prompt/decode) nao emite tool_call — a capacidade vem do difusor SFT.

Limitacao medida (documentada, nao escondida): abrindo o think do zero, a distribuicao inicial disputa We need (60%) vs Let me (40%) a temp 0.6, e o estilo minoritario (16/941 traces) deriva em verbosidade sem fechar o think em ate 6000 tokens. Dado um think no estilo dominante do corpus (426/941), o modelo fecha e emite tool_call em top-1 (</think> lp -0.48). Sondas em tmp/probe*.out; logs de geracao em tmp/infer_sft.log (dlora) e tmp/infer_base.log (base).

Files

File	What it is
difusor.ckpt	The SFT diffuser: 21 fp32 tensors, 22,032,736 parameters, step 237 (252.2 MiB, sha256 9023d59366b1cc42...). cfg inside records every hyperparameter below.
nanbeige_diffuser.gguf	The same diffuser packed as the sidecar GGUF for the llama.cpp dlora port (fixed seed-42 inference noise inside).
nanbeige_ffn_diffusion_lora.py	The adapter module: LoRADiffuser (DDIM), anexar_ffn_diffusion_lora, restaurar_rope (mandatory after from_pretrained), ruido_da_inferencia.
treinar_nanbeige_ffn_diffusion_lora.py	The trainer (transformers + Lightning Fabric, tyro CLI) that produced this checkpoint.
infer_nanbeige.py	The generation-verification recipe: NF4 load, attach, load delta, fixed noise, raw-trace prompt, sampled decode (temp 0.6, top_p 0.9) with the corpus-typical think appended as a style seed.

Inference (transformers)

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
from nanbeige_ffn_diffusion_lora import (
    anexar_ffn_diffusion_lora, restaurar_rope, ruido_da_inferencia)

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")
ck = torch.load("difusor.ckpt", map_location="cpu", weights_only=False)
lora_diffuser.load_state_dict(ck["lora_diffuser"])

# ruido fixo de inferencia (semente 42) -- sigma 1.0, igual ao GGUF publicado
estado = model.model.layers[0].mlp.estado
estado.ruido = ruido_da_inferencia(32, 1.0).to(model.device)

# prompt: template Qwen RAW (o SFT usou chat-format raw), terminando em
# <|im_start|>assistant\n -- NAO use apply_chat_template/add_generation_prompt
prompt += "</think>" + think_corpus   # semente de estilo (ver abaixo)

out = model.generate(input_ids=..., use_cache=False, max_new_tokens=6000,
                     do_sample=True, temperature=0.6, top_p=0.9,
                     eos_token_id=166101)  # <|im_end|>

use_cache=False is deliberate: the remote generate is incompatible with the installed transformers 5.x cache path (empty DynamicCache under num_loops>1, and a torch.min tuple bug in prepare_inputs_for_generation) — this is the same mode the trainer used.

Regime invariants

• Backbone in NF4 (load_in_4bit): the diffuser was trained against the bitsandbytes NF4 backbone; loading the original bf16 weights changes the backbone under the delta.
• Input x_attn (recorded in the checkpoint cfg).
• Frozen is not no_grad: gradients flow through the frozen backbone to the diffuser.
• restaurar_rope after every from_pretrained (transformers 5, remote code).
• Inference noise: one fixed draw, seed 42, sigma 1.0.
• 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
SFT data	tmp/traces.jsonl (raw chat-format, seq_len 9216)
SFT optimizer	lr 0.0002, warmup 10, 238 steps