inference-optimization / NemotronH-0.3B-A0.3B

NemotronH-0.3B-A0.3B

This is a tiny version of nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-Base-BF16 created for testing and development.

Model Details

• Base Model : nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-Base-BF16
• Architecture : nemotron_h (hybrid Mamba-2 / attention / mixture-of-experts)
• Total Parameters : 0.347B
• Activated Parameters : ~0.309B

The nemotron_h architecture is a hybrid that interleaves three block types. This tiny model preserves at least one of each so the full architecture is exercised:

• linear_attention — Mamba-2 SSM mixer
• full_attention — grouped-query attention mixer
• moe — mixture-of-experts FFN (non-gated experts + shared expert + latent projection)

The layer pattern is ["linear_attention", "moe", "linear_attention", "full_attention", "moe"] .

Configuration Changes

The following parameters were reduced from the original model:

Checkpoint Structure

Single-file checkpoint ( model.safetensors ). The tensor layout matches the original model: the backbone. prefix is used, and routed experts are stored as individual per-expert 2D tensors ( backbone.layers.N.mixer.experts.M.up_proj.weight / down_proj.weight ) rather than stacked 3D parameters. All non-MTP tensor names are identical (modulo layer/expert indices) to the original checkpoint. Multi-token-prediction ( mtp.* ) layers are intentionally omitted ( num_nextn_predict_layers = 0 ).

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("NemotronH-0.3B-A0.3B", device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("NemotronH-0.3B-A0.3B")

input_ids = tokenizer("According to all known laws", return_tensors="pt").input_ids.to(model.device)
output = model.generate(input_ids, max_new_tokens=20)
print(tokenizer.decode(output[0]))

Creation Process

This model was created using the llm-compressor create-tiny-model claude skill.

1. Reduced layers_block_type , hidden/attention/mamba sizes, and MoE sizes as above.
2. Re-initialized all weights, then fine-tuned on a small toy dataset until the perplexity target was reached (validates the model can learn).
3. Verified the saved checkpoint structure matches the original (excluding MTP layers).

Validation

Success: 1.0079255104064941 <= 10.0

==================================================
Generating sample text:
According to all known laws of aviation, there is no way a bee should be able to fly.
==================================================

Notes

• The Mamba mixer intermediate size is mamba_num_heads * mamba_head_dim (= 2048), and n_groups must divide it.
• use_mamba_kernels=True ; if the optional mamba_ssm / causal_conv1d packages are not installed, transformers falls back to a correct (slower) reference implementation.
• Because the large 131072-token vocabulary embedding dominates the parameter budget, total and activated parameter counts are both ~0.3B despite the MoE routing.