KissTheHabit/IDA-nemotron3-code-minitron-4b

IDA Nemotron-3 Code Minitron-4B

Code-focused SFT of nvidia/Minitron-4B-Base (NemotronForCausalLM architecture, Nemotron-3 family, 4.19B params, 32 layers, hidden=3072, relu2 activation), trained directly from the pristine pretrained checkpoint on a code-specific dataset (mined git-history diffs + a synthetic incorrect→correct→method corpus, plain-text formatted, no instruction-following data mixed in) — trained with NVFP4 tensor-core compute, published as a single-file MXFP8 checkpoint.

Why this is a separate repo: an earlier investigation (KissTheHabit/IDA-TRAIN-V2-nemotron-4b-mxfp8) found that fine-tuning this model on databricks/databricks-dolly-15k (general instruction-following text, no code) before touching code data consistently destroyed code-completion ability via catastrophic forgetting — every Dolly-first attempt scored pass@1≈0 on real MBXP regardless of SFT data format. This checkpoint tests training directly on code content from the pristine base, skipping that failure mode entirely. Note: the training dataset itself (mined from this project's own private git history) is not published here or anywhere — only the resulting model weights.

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).

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

model.safetensors 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.

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 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 — three different, independently true 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) — not a single tensor byte is touched, verified by a sha256 check on export. The declared shape stays the native engine's own [in_features, out_features] orientation (the transpose of what nn.Linear expects); that transpose happens only at load time, in memory, after decoding — never against the file on disk.

Status

Evaluated. Real, execution-based MBXP (amazon-science/mxeval, mbpp_release_v1, 974 Python problems, greedy decoding, pass@1). The weights currently in this repo are from the most recently listed experiment below; earlier rows are kept for a transparent record of what was tried.

Checkpoint	pass@1
this model	0.0
codeset_v4_expanded (37,618 rows, multi-language, 1 epoch)	0.002053388090349076

See the training repo's mlperf_log.txt MLLOG entries for full weight/compute precision provenance.

Loading

Standard transformers.AutoModelForCausalLM.from_pretrained does not understand the MXFP8 tensor pairs in this checkpoint directly. Build the model from config, decode with this project's own codec, and transpose the 2D projection weights:

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 = "."
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.