Nano-35B-A3B-128K

Nano-35B-A3B-128K

This is an architecture-only, untrained model targeting 131,072 tokens (128K). It uses ordinary pre-norm residual connections, interleaved local/global grouped-query attention, and DeepSeek-style MoE. There are no pretrained weights.

Architecture

Every layer operates on [batch, sequence, 3072]:

h = x + GQA(RMSNorm(x))
y = h + MoE(RMSNorm(h))

MoE(z) = SharedExpert(z) + sum(weight_i * RoutedExpert_i(z), i in top-2)

The two RMSNorm modules have separate weights. The model ends with a final RMSNorm and vocabulary projection. All 36 layers have independent weights; there is no recurrent core or weight sharing across layers. This candidate uses more layers and smaller experts to target approximately 36B total / 3B active parameters. Its quality and latency have not been benchmarked.

Local layers attend to the most recent 4,096 tokens. Global layers attend to every preceding token; they provide a direct path to earlier parts of the 128K sequence. The schedule is:

(local × 5, global) × 6

GQA shares each KV head across four query heads in both kinds of layer.

MoE and simplification

The MoE implementation remains DeepSeek V4's: normalized sqrtsoftplus routing scores, top-2 selection, routed output scale 1.5, routed SwiGLU pre-activation limit 10, and an additive shared expert. There is no gate multiplying the shared expert's output. All 36 layers use learned routing; the previous first-three-layer Hash-MoE and its token-routing tables are removed.

The local configuration derives directly from PreTrainedConfig. Obsolete hc_*, index_*, compress_*, attention LoRA, norm_topk_prob, num_nextn_predict_layers, and router_jitter_noise arguments are rejected, including loading overrides. Model construction revalidates the configuration before allocating weights.

What 128K means here

128K is a design and training target, not a demonstrated model capability. Setting the context limit and RoPE base does not teach long-document retrieval. The local/global ratio, window, and RoPE base still need long-context training and evaluation, including retrieval and tasks requiring information from multiple distant positions.

For batch size 1 and BF16 K/V, the logical persistent cache at 128K is about 2.601 GiB, compared with 13.500 GiB if all 36 layers used global GQA:

2 (K and V) × 6 KV heads × 128 dimensions × 2 bytes
× [6 × 131072 + 30 × (4096 - 1)]

The pinned dynamic cache retains window - 1 past tokens for each local layer. These numbers exclude weights, activations, allocator overhead, prefill temporaries, and backing storage retained by tensor views. They are not peak GPU-memory measurements.

With a window-aware kernel, attention work scales as O(30 × sequence_length × window + 6 × sequence_length²). The six global layers still have quadratic prefill cost. Local/global GQA does not retain V4's compressed million-token attention efficiency.

For long-context GPU runs, use a backend that implements the local window efficiently, such as a compatible FlashAttention build. The model forwards each layer's window explicitly to that backend. Eager/SDPA can construct dense local masks and should not be assumed memory-efficient for 128K prefill. The CPU validation below does not establish CUDA-kernel compatibility, 128K throughput, or full-size training feasibility.

Training defaults

The upstream initializer_range=0.02 is retained. Attention output and expert down projections have no additional depth scaling; training stability has not been measured.

output_router_logits defaults to false, avoiding retained per-layer router outputs and auxiliary-loss computation during inference. For training, explicitly pass output_router_logits=True together with labels to enable the inherited auxiliary balancing loss, with coefficient 0.001. This loss uses a softmax surrogate over router logits; expert selection uses the DeepSeek scoring function. The e_score_correction_bias buffer starts at zero and is not updated by an optimizer or automatic balancing controller here. This is not a reproduction of DeepSeek's complete training recipe.

For training, pass use_cache=False and output_router_logits=True. For padding-free packing, flatten documents into one sequence and reset position_ids at each document boundary. Omit attention_mask (or set it to None), even if it would contain only ones: an explicit 2D mask can bypass document-boundary detection in the pinned backends. With no mask or past KV cache, position IDs reach FlashAttention's varlen dispatch, while eager/SDPA build document-aware masks.

RMSNorm statistics are computed in FP32; norm weights follow the model dtype. Only the routing correction-bias buffer is kept in FP32 on reload.

Assisted generation remains disabled because sliding caches cannot roll back once old tokens have been discarded. Fullgraph compilation remains disabled pending separate validation.

The vocabulary size and BOS/EOS IDs remain those of the original architecture. The attention change does not require switching to a Qwen tokenizer; this repository does not package tokenizer files.

Initialization and accounting

Tested with Python 3.12, PyTorch 2.12.0, and Transformers 5.9.0. Install the pinned dependencies:

python -m pip install -r requirements.txt
python count_parameters.py
python init_empty_model.py

Parameter accounting includes the complete embedding and output matrices, all attention/norm/router weights, one shared expert, and two routed experts per layer for the active count. Active parameter counts are not FLOP counts.

Meta initialization avoids allocating the full model:

import torch
from transformers import AutoConfig, AutoModelForCausalLM

repo = "bowang0911/Nano-35B-A3B-128K"
config = AutoConfig.from_pretrained(repo, trust_remote_code=True)
with torch.device("meta"):
    model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)

Use a reviewed revision when loading remote code. from_config creates random weights; from_pretrained cannot load weights from this architecture-only repository. The v4_nano model type needs explicit support in inference engines. Earlier V4 attention checkpoints are incompatible with this revision.

Validation

python -m unittest -v test_model

Fourteen CPU tests exercise a real tiny model with local/global GQA and DeepSeek MoE: all parameter tensors receive finite nonzero training gradients, single-stream residual equations, causal/local/global attention masks, cached decoding and chunked prefill across window boundaries, grouped KV cache sizes, left padding, high position IDs, checkpointed gradients, and FP32/BF16 save/reload/generation through AutoModel. They also cover packed document isolation, invalid loading overrides, obsolete options, and opt-in router outputs and balancing loss. A meta-device check compares the full model's actual parameter count with the analytical count.

Model-level FlashAttention tests verify that packed position IDs and per-layer windows reach its low-level helper, and that the real dispatch logic selects the varlen path with separate document boundaries. External kernels are stubbed; these are not GPU FlashAttention execution tests. High position IDs are tested on a short sequence, not a complete 128K sequence.

References and license

Local model wiring follows the Apache-2.0 Transformers implementation; see LICENSE-APACHE-2.0.

• Transformers DeepSeek V4 documentation
• Transformers Llama documentation
• Gemma 3 local/global attention design demonstrates this attention pattern at 128K in a separately trained model. This configuration uses different dimensions, MoE, normalization, and window sizes.