Models / nemotron_labs_diffusion

Nemotron-Labs-Diffusion-8B: GPU memory to fine-tune

Fine-tuning Nemotron-Labs-Diffusion-8B at batch 1 and 2048 tokens takes about 160 GB of GPU memory for a full fine-tune (2 x H200 141 GB, split by layer), and 25.5 GB with LoRA r=16 on a bf16 base (borderline on one A100 40 GB). Estimated from the model's structure, reading the config under transformers semantics, not measured. What fits on which card.

Reconstructed from its own config.json with no weights read. 107K downloads on Hugging Face.

Our count against the checkpoint

The left number comes from the graph. The right one is the number of scalars in the published weight files. Nothing on this page was tuned to make them agree.

Derived from structure

7.89B

7,886,753,792 parameters

In the published checkpoint

8.49B

8,489,553,920 scalars · safetensors.total, read 2026-06-03

Delta

-7.10%

custom-code This repository ships its own modeling code (`auto_map`, e.g. `configuration_nemotron_labs_diffusion.py`), so `config.json` names a class in the repo rather than an architecture `transformers` defines. The graph below is what those config keys mean under `transformers` semantics, which is not necessarily what the repo's own file builds. A gap here is a statement about what we read, not about the checkpoint.

What it costs to run

Cost is a roofline estimate on the priced GPU for 10 epochs at batch 32, 2048 tokens over 50,000 samples (assumed; no dataset attached).

Layers

138

Will it forward-pass

Yes

Priced on

A10G (24GB)

Est. one run

$502.88

Memory: what fits on which card

Largest batch per card. Training is sized at 2048 tokens per sample, serving at 4096 tokens per sequence.

Full fine-tune, batch 1 at 2048 tokens

160 GB

LoRA r=16, batch 1 at 2048 tokens

25.5 GB

Serve: bf16 weights

15.8 GB

Serve: KV cache, 4096 tokens

2.28 GB

Card	Full fine-tune	LoRA r=16	Serve bf16, 4096 tokens
T4 16 GB	does not fit at batch 1	does not fit at batch 1	does not fit at batch 1
L4 24 GB	does not fit at batch 1	does not fit at batch 1	batch 2
A100 40 GB	does not fit at batch 1	batch 2, borderline	batch 8
A100 / H100 80 GB	does not fit at batch 1	batch 5	batch 23
H200 141 GB	2 cards at batch 1 (98 + 40 layers)	batch 11	batch 46

Training: 2048 tokens per sample (2048, or the model's context when shorter; the structure table is at the config's maximum, 262144), AdamW with bf16 autocast, SDPA attention (no score matrix stored), 1.15x headroom on every card. Full fine-tune holds weights, gradients and optimizer state (1.2 x 16 bytes per parameter) plus activations. LoRA is r=16 on every linear layer but the head, base frozen in bf16 (2 bytes per parameter), adapters at full AdamW state, same activations. A two-card split is a contiguous layer cut (pipeline parallel), computed only where one card cannot hold batch 1. "Borderline" means the card holds batch 1 on the estimate but not on the conservative peak (the estimate x 1.43, at or above the measured peak on 95% of held-out runs): measure one step before reserving it. Serving: bf16 weights plus the KV cache of 4096 tokens per sequence; prefill working set, CUDA context and framework reserve not counted.

These are estimates from the graph, not measurements. Against measured runs, held out by architecture, the estimate flagged 113 of the 131 jobs that did not fit their card and said "fits" on the other 18. Every measured run.

Structure

140 nodes. Output shapes are propagated from the input shape, batch dimension excluded, at 262144 tokens: the config's maximum context, not the 2048 the memory above is sized at.

	Layer	Type	Output shape
1	Input	Input	1 × 262144
2	Embedding	Embedding	1 × 262144 × 4096
3	Positional_Embedding	Learned Pos Embed	1 × 262144 × 4096
4	Attention_1	Multi-Head Attention	1 × 262144 × 4096
5	Add_1	Add	1 × 262144 × 4096
6	LayerNorm_1_1	LayerNorm	1 × 262144 × 4096
7	FFN_1	Feed Forward	1 × 262144 × 4096
8	Attention_2	Multi-Head Attention	1 × 262144 × 4096
9	Add_2	Add	1 × 262144 × 4096
10	LayerNorm_2_1	LayerNorm	1 × 262144 × 4096
11	FFN_2	Feed Forward	1 × 262144 × 4096
12	Attention_3	Multi-Head Attention	1 × 262144 × 4096
13	Add_3	Add	1 × 262144 × 4096
14	LayerNorm_3_1	LayerNorm	1 × 262144 × 4096
15	FFN_3	Feed Forward	1 × 262144 × 4096
16	Attention_4	Multi-Head Attention	1 × 262144 × 4096
17	Add_4	Add	1 × 262144 × 4096
18	LayerNorm_4_1	LayerNorm	1 × 262144 × 4096
19	FFN_4	Feed Forward	1 × 262144 × 4096
20	Attention_5	Multi-Head Attention	1 × 262144 × 4096
21	Add_5	Add	1 × 262144 × 4096
22	LayerNorm_5_1	LayerNorm	1 × 262144 × 4096
23	FFN_5	Feed Forward	1 × 262144 × 4096
24	Attention_6	Multi-Head Attention	1 × 262144 × 4096
25	Add_6	Add	1 × 262144 × 4096
26	LayerNorm_6_1	LayerNorm	1 × 262144 × 4096
27	FFN_6	Feed Forward	1 × 262144 × 4096
28	Attention_7	Multi-Head Attention	1 × 262144 × 4096
29	Add_7	Add	1 × 262144 × 4096
30	LayerNorm_7_1	LayerNorm	1 × 262144 × 4096
31	FFN_7	Feed Forward	1 × 262144 × 4096
32	Attention_8	Multi-Head Attention	1 × 262144 × 4096
33	Add_8	Add	1 × 262144 × 4096
34	LayerNorm_8_1	LayerNorm	1 × 262144 × 4096
35	FFN_8	Feed Forward	1 × 262144 × 4096
36	Attention_9	Multi-Head Attention	1 × 262144 × 4096
37	Add_9	Add	1 × 262144 × 4096
38	LayerNorm_9_1	LayerNorm	1 × 262144 × 4096
39	FFN_9	Feed Forward	1 × 262144 × 4096
40	Attention_10	Multi-Head Attention	1 × 262144 × 4096
41	Add_10	Add	1 × 262144 × 4096
42	LayerNorm_10_1	LayerNorm	1 × 262144 × 4096
43	FFN_10	Feed Forward	1 × 262144 × 4096
44	Attention_11	Multi-Head Attention	1 × 262144 × 4096
45	Add_11	Add	1 × 262144 × 4096
46	LayerNorm_11_1	LayerNorm	1 × 262144 × 4096
47	FFN_11	Feed Forward	1 × 262144 × 4096
48	Attention_12	Multi-Head Attention	1 × 262144 × 4096
49	Add_12	Add	1 × 262144 × 4096
50	LayerNorm_12_1	LayerNorm	1 × 262144 × 4096
51	FFN_12	Feed Forward	1 × 262144 × 4096
52	Attention_13	Multi-Head Attention	1 × 262144 × 4096
53	Add_13	Add	1 × 262144 × 4096
54	LayerNorm_13_1	LayerNorm	1 × 262144 × 4096
55	FFN_13	Feed Forward	1 × 262144 × 4096
56	Attention_14	Multi-Head Attention	1 × 262144 × 4096
57	Add_14	Add	1 × 262144 × 4096
58	LayerNorm_14_1	LayerNorm	1 × 262144 × 4096
59	FFN_14	Feed Forward	1 × 262144 × 4096
60	Attention_15	Multi-Head Attention	1 × 262144 × 4096
61	Add_15	Add	1 × 262144 × 4096
62	LayerNorm_15_1	LayerNorm	1 × 262144 × 4096
63	FFN_15	Feed Forward	1 × 262144 × 4096
64	Attention_16	Multi-Head Attention	1 × 262144 × 4096
65	Add_16	Add	1 × 262144 × 4096
66	LayerNorm_16_1	LayerNorm	1 × 262144 × 4096
67	FFN_16	Feed Forward	1 × 262144 × 4096
68	Attention_17	Multi-Head Attention	1 × 262144 × 4096
69	Add_17	Add	1 × 262144 × 4096
70	LayerNorm_17_1	LayerNorm	1 × 262144 × 4096
71	FFN_17	Feed Forward	1 × 262144 × 4096
72	Attention_18	Multi-Head Attention	1 × 262144 × 4096
73	Add_18	Add	1 × 262144 × 4096
74	LayerNorm_18_1	LayerNorm	1 × 262144 × 4096
75	FFN_18	Feed Forward	1 × 262144 × 4096
76	Attention_19	Multi-Head Attention	1 × 262144 × 4096
77	Add_19	Add	1 × 262144 × 4096
78	LayerNorm_19_1	LayerNorm	1 × 262144 × 4096
79	FFN_19	Feed Forward	1 × 262144 × 4096
80	Attention_20	Multi-Head Attention	1 × 262144 × 4096
81	Add_20	Add	1 × 262144 × 4096
82	LayerNorm_20_1	LayerNorm	1 × 262144 × 4096
83	FFN_20	Feed Forward	1 × 262144 × 4096
84	Attention_21	Multi-Head Attention	1 × 262144 × 4096
85	Add_21	Add	1 × 262144 × 4096
86	LayerNorm_21_1	LayerNorm	1 × 262144 × 4096
87	FFN_21	Feed Forward	1 × 262144 × 4096
88	Attention_22	Multi-Head Attention	1 × 262144 × 4096
89	Add_22	Add	1 × 262144 × 4096
90	LayerNorm_22_1	LayerNorm	1 × 262144 × 4096
91	FFN_22	Feed Forward	1 × 262144 × 4096
92	Attention_23	Multi-Head Attention	1 × 262144 × 4096
93	Add_23	Add	1 × 262144 × 4096
94	LayerNorm_23_1	LayerNorm	1 × 262144 × 4096
95	FFN_23	Feed Forward	1 × 262144 × 4096
96	Attention_24	Multi-Head Attention	1 × 262144 × 4096
97	Add_24	Add	1 × 262144 × 4096
98	LayerNorm_24_1	LayerNorm	1 × 262144 × 4096
99	FFN_24	Feed Forward	1 × 262144 × 4096
100	Attention_25	Multi-Head Attention	1 × 262144 × 4096
101	Add_25	Add	1 × 262144 × 4096
102	LayerNorm_25_1	LayerNorm	1 × 262144 × 4096
103	FFN_25	Feed Forward	1 × 262144 × 4096
104	Attention_26	Multi-Head Attention	1 × 262144 × 4096
105	Add_26	Add	1 × 262144 × 4096
106	LayerNorm_26_1	LayerNorm	1 × 262144 × 4096
107	FFN_26	Feed Forward	1 × 262144 × 4096
108	Attention_27	Multi-Head Attention	1 × 262144 × 4096
109	Add_27	Add	1 × 262144 × 4096
110	LayerNorm_27_1	LayerNorm	1 × 262144 × 4096
111	FFN_27	Feed Forward	1 × 262144 × 4096
112	Attention_28	Multi-Head Attention	1 × 262144 × 4096
113	Add_28	Add	1 × 262144 × 4096
114	LayerNorm_28_1	LayerNorm	1 × 262144 × 4096
115	FFN_28	Feed Forward	1 × 262144 × 4096
116	Attention_29	Multi-Head Attention	1 × 262144 × 4096
117	Add_29	Add	1 × 262144 × 4096
118	LayerNorm_29_1	LayerNorm	1 × 262144 × 4096
119	FFN_29	Feed Forward	1 × 262144 × 4096
120	Attention_30	Multi-Head Attention	1 × 262144 × 4096
121	Add_30	Add	1 × 262144 × 4096
122	LayerNorm_30_1	LayerNorm	1 × 262144 × 4096
123	FFN_30	Feed Forward	1 × 262144 × 4096
124	Attention_31	Multi-Head Attention	1 × 262144 × 4096
125	Add_31	Add	1 × 262144 × 4096
126	LayerNorm_31_1	LayerNorm	1 × 262144 × 4096
127	FFN_31	Feed Forward	1 × 262144 × 4096
128	Attention_32	Multi-Head Attention	1 × 262144 × 4096
129	Add_32	Add	1 × 262144 × 4096
130	LayerNorm_32_1	LayerNorm	1 × 262144 × 4096
131	FFN_32	Feed Forward	1 × 262144 × 4096
132	Attention_33	Multi-Head Attention	1 × 262144 × 4096
133	Add_33	Add	1 × 262144 × 4096
134	LayerNorm_33_1	LayerNorm	1 × 262144 × 4096
135	FFN_33	Feed Forward	1 × 262144 × 4096
136	Attention_34	Multi-Head Attention	1 × 262144 × 4096
137	Add_34	Add	1 × 262144 × 4096
138	LayerNorm_34_1	LayerNorm	1 × 262144 × 4096
139	FFN_34	Feed Forward	1 × 262144 × 4096
140	Output	Output	1 × 262144 × 4096

What the verifier says

info34 attention layers at embedDim 4096 cache full per-head K/V: about 544 KB per token at fp16, which dominates memory at long context. Grouped-query attention (e.g. 8:1) would cut this ~8×; multi-head latent attention (MLA) shrinks it ~10× or more. This is the move production LLMs make; it does not change the parameter count. Fix: Switch attention to groupedQueryAttention (set numKVHeads below numHeads, e.g. numHeads/4) or mla (a low-rank cached latent).
full-mha-serving-cost

infoAt 34 stacked attention layers, residual-branch outputs add up; unscaled init lets activation variance grow with depth. GPT-2/LLaMA-family models scale the residual projections by depth (N(0, 0.02 / √(2L))). Fix: Scale residual output projections by depth: nn.init.normal_(w, std=0.02 / math.sqrt(2 * n_layers))
deep-attention-default-init

warnAcross 34 attention layers this design caches 544 KB per token, so a single 8,192-token sequence needs ~4.6 GB of KV cache before weights or activations. That exceeds the 4 GB budget this rule assumes for serving headroom. Fix: Cut KV width: raise the GQA ratio (fewer numKVHeads), switch to MLA, reduce depth or embedDim, or accept a shorter serving context.
kv-cache-context-budget

What changed upstream

The change tracker follows nvidia/Nemotron-Labs-Diffusion-8B: every revision it has read, what changed in the config, tokenizer, chat template or weights between them, and an OSV advisory per change.

Revision history and changes Badge JSON OSV advisories All tracked changes

Do this to your own model

Same numbers, on a model in your repo, in one command. No account.

pip install neurarch-trace
neurarch-trace nvidia/Nemotron-Labs-Diffusion-8B --plan --share

Paste a model instead Markdown twin Index JSON POST /api/v1/plan