nemo_automodel.components.models.mistral_large4.flops

Count useful GEMM FLOPs for text-only full-parameter training of Mistral Large 4.

Trainable projections and attention count forward, input gradients and weight gradients (six FLOPs per MAC). The DSA indexer is a frozen hard selector, so it runs only on full layers and counts forward only (two FLOPs per MAC): its projections plus scores over the causal keys. Per layer:

• q_proj: hidden * heads * head_dim (no Q low-rank path).
• kv_a: hidden * (kv_lora_rank + rope_dim) on layers that own their KV only; NoPE layers have no rope part and layers that share KV have no kv_a.
• kv_b: kv_lora_rank * heads * (nope_dim + v_head_dim) on every layer, with nope_dim = head_dim on NoPE layers and qk_nope_head_dim otherwise.
• o_proj: heads * v_head_dim * hidden.
• Attention: pairs * heads * (head_dim + v_head_dim) MACs, where a query at position p attends min(p + 1, index_topk) keys (exact, not S * topk).
• FFN: layers below first_k_dense_replace are dense SwiGLU; the others run num_experts_per_tok routed experts and n_shared_experts shared experts of moe_intermediate_size plus the router.
• LM head: hidden * vocab (the embedding lookup has no GEMM).

Counted work is the algorithmic model FLOPs (MFU), independent of the attention backend: the absorbed TileLang kernel and dense SDPA execute different amounts of work. Activation recomputation, padding, elementwise operations (norms, SSMax, LeRoPE, softcap, softmax, top-k), optimizer updates and communication are excluded. Vision and the draft head are not executed by the causal LM and are excluded.

Each sequence is one document without padding. Packed-document FLOPs need the individual document lengths and cannot be inferred from a packed row length.

Parameters:

Returns: float

Useful model FLOPs over all devices for one optimizer step.

Raises:

• ValueError: gbs or seq_len is not positive.