notpaulmartin/SmolLM3-3B-flax

SmolLM3-3B Flax

A BF16 Flax conversion of HuggingFaceTB/SmolLM3-3B, for inference and LoRA training with Tokenkit. No additional training or quantization was applied.

Source revision: a07cc9a04f16550a088caea529712d1d335b0ac1.

The checkpoint contains 3,075,098,624 parameters in 326 tensors. It preserves the original weights, tokenizer, chat template, generation settings, and tied embeddings. Linear weights are transposed into Flax's kernel layout. The model uses Tokenkit's tpu_llama implementation, with RoPE disabled at zero-based layers 3, 7, ..., 35. In no_rope_layers, 1 means RoPE and 0 means NoPE, matching the upstream implementation. The legacy pretraining_tp field is normalized to 1 because native SmolLM3 does not use it.

Tokenkit

Use a Tokenkit checkout with SmolLM3 support:

model:
  pretrained_model_name_or_path: notpaulmartin/SmolLM3-3B-flax
  tokenizer_name: notpaulmartin/SmolLM3-3B-flax:source=SmolLM3

The existing LoRA trainer and generator work with this checkpoint. The tokenizer kind maps canonical chat markers to ChatML and uses <|im_end|> as EOS and padding, without inserting BOS. Use the original apply_chat_template when the model's system metadata and reasoning-mode instructions are required. Canonical ChatML formatting alone does not add them.

Standalone Flax loading

The repository includes the configuration and Flax model code. Review it before enabling trust_remote_code. This example runs a forward pass on one JAX device; set JAX_PLATFORMS=cpu before Python for CPU execution.

import jax
import jax.numpy as jnp
import numpy as np
from transformers import AutoConfig, AutoTokenizer, FlaxAutoModelForCausalLM

repo = "notpaulmartin/SmolLM3-3B-flax"
config = AutoConfig.from_pretrained(repo, trust_remote_code=True)
config.mesh = jax.sharding.Mesh(
    np.array(jax.devices()[:1]).reshape(1, 1, 1), ("data", "fsdp", "model")
)
config.max_length = 128
model, params = FlaxAutoModelForCausalLM.from_pretrained(
    repo, config=config, trust_remote_code=True,
    _do_init=False, dtype=jnp.float32,
)
tokenizer = AutoTokenizer.from_pretrained(repo)
inputs = tokenizer("The capital of France is", return_tensors="np")
logits = model(inputs["input_ids"], params=params).logits
print(tokenizer.decode([int(logits[0, -1].argmax())]))

The mesh and max_length settings are required at runtime. This TPU-oriented implementation uses Tokenkit's generation utilities instead of Transformers' .generate() method. It does not require a TPU for eager-attention forward passes.

Validation

Tested with Transformers 4.57.1, JAX 0.8.2, and Flax 0.12.2 on CPU:

• The full converted checkpoint matches upstream PyTorch logits on a 78-token chat prompt. Both models used FP32 computation over the original BF16 weights. Maximum absolute error was 0.00024748; mean absolute error was 0.00001636. This is numerical parity, not bitwise BF16-compute equivalence.
• Tokenkit's cached generator answered "The capital of France is Paris."
• Small-model tests cover upstream parity, scanned and unscanned execution, cached decoding, gradient checkpointing, a loss-reducing LoRA update, runtime-versus-merged LoRA, and BF16 conversion/tokenizer roundtrip.

See validation.json for full-checkpoint metrics. Long-context behavior and full TPU training runs were not benchmarked. The original 65,536-position configuration is preserved; this conversion does not add YaRN extrapolation.

The original Hugging Face team trained the model. Its capabilities, limitations, and license apply unchanged. See the original model card and the included ORIGINAL_MODEL_CARD.md.