KissTheHabit/IDA-TRAIN-V2-nemotron-4b-mxfp8

IDA-TRAIN-V2 Nemotron/Minitron-4B, MXFP8

Checkpoints from an ongoing investigation into SFT-ing NemotronForCausalLM (4.19B params, 32 layers, hidden=3072, relu2 activation) and scoring it with real, execution-based MBXP (mbpp) pass@1 — trained with NVFP4 tensor-core compute, published as single-file MXFP8 checkpoints. This repo documents three failed instruction-tuning attempts and why, plus the experiment currently testing the actual fix.

Hardware

Trained on 2x consumer-Blackwell RTX 5070 Ti (sm_120a, 16GB each) through this project's native C++/CUDA engine, model-parallel across both cards (host-staged 1F1B pipeline — no CUDA P2P on consumer Blackwell). Eval in this repo runs the same two-card split via accelerate's device map.

The investigation so far

Three separate SFT attempts on databricks/databricks-dolly-15k (general instruction-following data, no code) all scored pass@1 ≈ 0 with repetition-collapse output — the model gets stuck greedily repeating a token or whitespace character instead of writing code:

Attempt	Base	Data format	pass@1	Failure signature
dolly15k_sft/	nvidia/Minitron-4B-Instruct	Alpaca-style ### Instruction:/### Response:	~0 (see note)	repetition collapse (descendant codeset_v3_sft/ confirmed 0.0 with the same pattern)
dolly15k_base_sft/	nvidia/Minitron-4B-Base	same ###-marked template	0.0	word-repeat: "are are are are...", "isisisisisis..."
dolly15k_plain_sft/	nvidia/Minitron-4B-Base	plain flowing text, no markers	0.0	whitespace collapse: output is empty or hundreds of repeated tabs

Each hypothesis was tested and eliminated in turn, not assumed:

1. "It's the Instruct model's missing chat template." Ruled out — the untouched, non-instruct nvidia/Minitron-4B-Base checkpoint produces coherent code on the exact same raw-completion prompt via plain transformers.AutoModelForCausalLM.from_pretrained, with zero custom pipeline code involved.
2. "It's the ###-marked SFT template teaching a rigid turn structure." Ruled out — reformatting Dolly to plain flowing text (no markers at all) and retraining produced the same pass@1=0.0, just a different degenerate output pattern.
3. "It's a bug in this repo's checkpoint conversion/loading pipeline" (the MXFP8 decode, the native→HF orientation transpose, the LayerNorm1p fold/unfold). Ruled out — a checkpoint trained for only 4 optimizer steps from the pristine Base model, run through the exact same production pipeline (native_checkpoint_to_hf.py → gen_mbxp_completions_mxfp8.py), produced coherent Python. The pipeline is verified correct.
4. Current leading hypothesis: Dolly's content, not its format, is the problem. Dolly-15k is 100% non-code natural-language data (recipes, trivia, general Q&A). Two full epochs of full-parameter fine-tuning on that at a real learning rate can displace code-completion capability through ordinary catastrophic forgetting — no bug required. codeset_v3 (this repo's own code-specific dataset) trained directly from the pristine Base checkpoint, skipping Dolly entirely, is the live test of this — see the table below for status.

Note on dolly15k_sft/'s own pass@1: its individual score was not preserved (lost when the original training host went down mid-session and had to be rebuilt from this repo's own published weights). Its direct descendant codeset_v3_sft/, continued from it, scored 0.0 with the identical repetition-collapse signature, which is why it's listed as ~0 rather than left blank.

Precision recipe

• Compute: NVFP4 tensor-core matrix multiply (native_mma_sm120a)
• Weight master: MXFP8 (mxfp8_e4m3_ue8m0_k32 — E4M3 payload, one UE8M0 power-of-two scale per 32 input values), device-resident
• Gradients / Lion optimizer arithmetic: BF16

Each model.safetensors in this repo carries that MXFP8 weight master straight through: attention and MLP projection weights (q/k/v/o_proj, up/down_proj) are stored as real F8_E4M3 payload + UE8M0 scale tensor pairs (__metadata__.weights_dtype == "MXFP8_E4M3_UE8M0_K32") in the same file as the plain-BF16 embeddings/norms/biases — one checkpoint, not a separate quantized bundle alongside a full-precision copy.

What this is not: a live NVFP4 tensor-core inference checkpoint. Nothing in this project currently serves generation through NVFP4/FP8 tensor cores — that path is training-only. Loading these weights for inference (this repo's own eval included) decodes the MXFP8 payload back to BF16 and runs standard F.linear. So: NVFP4 compute at training time, MXFP8 weights at rest, BF16 compute at inference time — each claim is real and independently true, but they are three different statements, not one.

The published model.safetensors is byte-identical to the trainer's persisted master. Only the JSON header's tensor names are rewritten (to transformers' Llama-family module convention, e.g. model.layers.{i}.self_attn.q_proj.weight) — not a single tensor byte is touched, verified by a sha256 check on every export. The declared shape stays the native engine's own [in_features, out_features] orientation, which is the transpose of what NemotronForCausalLM's nn.Linear layers expect ([out_features, in_features]); MXFP8 blocks group 32 elements along the last axis, and that axis differs before and after a transpose, so a byte-identical file that's also already in HF orientation isn't possible. The transpose therefore happens only at load time, in memory, after decoding — never against the file on disk. See Loading below.

Checkpoints in this repo

Path	Base	Data	Status	pass@1
dolly15k_sft/	nvidia/Minitron-4B-Instruct	Dolly-15k, ###-marked	broken (repetition collapse)	~0
codeset_v3_sft/	continues from dolly15k_sft/	this repo's git history + synthetic method corpus, ###-marked	broken (repetition collapse)	0.0
dolly15k_base_sft/	nvidia/Minitron-4B-Base	Dolly-15k, ###-marked	broken (repetition collapse)	0.0
dolly15k_plain_sft/	nvidia/Minitron-4B-Base	Dolly-15k, plain text	broken (repetition collapse)	0.0
codeset_v3_genesis_plain_sft/	nvidia/Minitron-4B-Base (direct — Dolly skipped)	this repo's git history + synthetic method corpus, plain text	training	(pending)

Evaluation

Real, execution-based MBXP (amazon-science/mxeval, mbpp_release_v1, 974 Python problems, greedy decoding, pass@1) — not a synthetic or held-out split.

Checkpoint	pass@1	weight precision
baseline (pre-SFT, nvidia/Minitron-4B-Instruct)	0.0298	bf16 (never quantized)
dolly15k_sft/	~0 (see investigation notes)	mxfp8_e4m3_ue8m0_k32
codeset_v3_sft/	0.0	mxfp8_e4m3_ue8m0_k32
dolly15k_base_sft/	0.0	mxfp8_e4m3_ue8m0_k32
dolly15k_plain_sft/	0.0	mxfp8_e4m3_ue8m0_k32
codeset_v3_genesis_plain_sft/	(training)	mxfp8_e4m3_ue8m0_k32

Each checkpoint's own real/measured eval_accuracy, weight_precision, and compute_precision provenance is also recorded as MLLOG entries in that run's mlperf_log.txt on the training box.

Loading

Standard transformers.AutoModelForCausalLM.from_pretrained does not understand the MXFP8 tensor pairs in this checkpoint directly (F8_E4M3 + .scale companions aren't a transformers-native format), regardless of orientation. Build the model from config, decode with this project's own codec, and transpose the 2D projection weights (the file itself stays native-oriented — see above):

from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
import torch, sys
sys.path.insert(0, "path/to/IDA-TRAIN-V2/src")
from ida_train.native.mxfp8_codec import read_mxfp8_safetensors

model_dir = "dolly15k_plain_sft"  # any subfolder above
tokenizer = AutoTokenizer.from_pretrained(model_dir)
config = AutoConfig.from_pretrained(model_dir)
model = AutoModelForCausalLM.from_config(config, torch_dtype=torch.bfloat16)

decoded = read_mxfp8_safetensors(f"{model_dir}/model.safetensors")
proj_leaves = (".self_attn.q_proj.weight", ".self_attn.k_proj.weight",
               ".self_attn.v_proj.weight", ".self_attn.o_proj.weight",
               ".mlp.up_proj.weight", ".mlp.down_proj.weight", ".mlp.gate_proj.weight")
for name, arr in decoded.items():
    if arr.ndim == 2 and name.endswith(proj_leaves):
        decoded[name] = arr.T.copy()  # native [in,out] -> HF nn.Linear [out,in]

state_dict = {k: torch.from_numpy(v).to(torch.bfloat16) for k, v in decoded.items()}
model.load_state_dict(state_dict, strict=True)

A working reference implementation (plus GPU dispatch, batching, and stop sequences for generation) is gen_mbxp_completions_mxfp8.py in the training repo.

Given every checkpoint above except codeset_v3_genesis_plain_sft/ is confirmed broken, none of them are currently recommended for actual use. This repo is published as a transparent, real-evidence record of the debugging process, not as a working model release.