09505fcc6e* Update VRAM estimator to cater to broader model configs * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix attn backend check, better support for MoE etc * Studio: tighten VRAM estimator structured-shape and attention paths - Conservative attention fallback: when resolve_attention_implementation fails, charge the quadratic non-flash activation path instead of silently keeping the optimistic flash_attention_2 default. - Resolve attention on a shallow config copy so _set_attn_impl does not mutate the cached config returned by _load_config_for_gpu_estimate. - Use getattr for AutoModelForCausalLM._model_mapping to avoid raising on private-attribute renames in transformers. - Treat sdpa as O(n) linear attention; PyTorch SDPA dispatches to flash or memory-efficient backends, only eager needs the quadratic term. - Per-layer activation accounting: structured archs (head_dim, layer_types, attention_k_eq_v, num_kv_shared_layers, double-wide MLP) now flow into compute_activation_bytes via _text_linear_dims, instead of using the legacy hidden_size//num_attention_heads KV/MLP shape. - Exclude MLA configs (q_lora_rank set) from the structured-shape path so q_lora low-rank projection formulas keep applying when head_dim is also present. - _build_text_module_elements emits a single MLA self_attn aggregate using _compute_attn_elements when q_lora_rank is set, avoiding the ~10% overcount that fed into _compute_skipped_quantizable_elements. - Restrict _module_path_matches to known text-tower prefixes so VLM skip names like vision_tower.model.layers.<i>.self_attn.q_proj no longer falsely shadow the text alias model.layers.<i>.self_attn.q_proj. - Pick up enable_moe_block from the config and add the per-layer dense MLP alongside the MoE experts in compute_total_params and compute_lora_params (Gemma4-style parallel dense + MoE block). - Single-pass structured layer accounting in _compute_layer_elements, removing the duplicate _text_linear_dims walks. - Drop the now-zero (activations - activations_computed) shard term in VramBreakdown.min_gpu_vram and the stale comment that referred to it. - attention_implementation typed as Optional[str] to match call sites that pass None. - Inline rationale comments on DOUBLE_QUANT_4BIT_FACTOR and NON_FLASH_ATTENTION_FACTOR pointing at VRAM_ESTIMATION.md. * Studio: extend parallel-MoE accounting + non-prefix dense layer support - Apply enable_moe_block / moe_has_dense_mlp symmetrically: activation per-layer MLP size in _layer_qkv_mlp_sizes now adds the parallel dense MLP for MoE layers, matching the weight and LoRA accounting added in the prior commit. Skip-quantizable mapping in _build_text_module_elements now registers both mlp.experts and per-projection mlp.{name} entries for MoE layers when the parallel dense block is present, so an llm_int8_skip_modules entry like "model.layers.N.mlp" covers both. - Track dense layer indices as a tuple (dense_layer_indices) extracted from first_k_dense_replace or decoder_sparse_step + mlp_only_layers, and dispatch dense-vs-MoE accounting through _is_dense_mlp_layer. The prior count-based path silently mis-bucketed layers when mlp_only_layers was non-prefix (e.g. [3, 5] on an 8-layer model). num_dense_layers is derived from len(dense_layer_indices) for backward compatibility. - Drop the redundant ">0" check in _is_kv_shared_layer so configs with num_kv_shared_layers == num_hidden_layers (every layer shared) are correctly recognized as shared. - Refresh VRAM_ESTIMATION.md section 5 to note that sdpa joins flash_attention_2 in the linear activation path; refresh the VramBreakdown.activations_computed comment now that the activation floor is gone. * Studio: Gemma4 PLE accounting, flex_attention, KV-share guard restore - Add flex_attention to LINEAR_ATTENTION_IMPLS. Unsloth's resolve_attention_implementation returns "flex_attention" when HAS_FLASH_ATTENTION is False and the model class supports flex; PyTorch FlexAttention is a memory-efficient kernel, not a quadratic eager attention path. Without this, activation estimates over-charge ~36x. - Restore the `> 0` guard in _is_kv_shared_layer. Transformers Gemma4 (modeling_gemma4.py:1031, modular_gemma4.py:863, :926) uses `layer_idx >= first_kv_shared_layer_idx > 0`, so configs that mark every layer as KV-shared raise on construction. Reverting the unconditional acceptance avoids producing a detailed estimate for a shape the actual model code rejects. - Extend the parallel dense MLP path (`enable_moe_block`) in _build_text_module_elements: when the arch is non-structured, use arch.intermediate_size for the dense gate/up/down dims instead of _text_linear_dims (which returns moe_intermediate_size via _get_mlp_size). Prior code under-counted skipped quantizable elements for the parallel dense block by up to 8x on GLM-style configs. - Add Gemma4 per-layer-input (PLE) module accounting: per_layer_model_projection (one global Linear) plus per-layer per_layer_input_gate and per_layer_projection are added to the quantizable text-linear total in _compute_layer_elements; post_per_layer_input_norm and per_layer_projection_norm flow into the non-quantizable bucket. compute_lora_params adds the same three Linear modules to the all-linear total. References: transformers_versions/5.7.0/.../gemma4/modular_gemma4.py:1077-1083, :1247-1253. - VRAM_ESTIMATION.md section 5 now lists flex_attention alongside sdpa and flash_attention_2 as linear-memory backends. * Studio: shared-expert variants, mlp_layer_types dispatch, PLE skip, all-linear str, deepcopy resolver Five targeted estimator corrections: - _compute_dense_layer_indices now reads `mlp_layer_types` ahead of `first_k_dense_replace` / `decoder_sparse_step`. Transformers Exaone-MoE, Laguna, Hy_v3, GLM-MoE-DSA, GLM4-MoE-Lite, Ernie4_5_VL_MoE etc. ship the per-position list and may omit the prefix-style fields entirely. - _build_text_module_elements registers per_layer_input_gate / per_layer_projection (per layer) and per_layer_model_projection (global) in the canonical element map and alias map. The PLE element count was added to total_quantizable in a prior commit but skip-module matching against names like model.layers.0.per_layer_input_gate produced 0-byte delta. Layer aggregate text.layers.<i> now sums all layer modules so prefix skip names cover the PLE pieces too. - _targets_all_linear coerces a bare string `"all-linear"` to `["all-linear"]` before set comparison; the previous set comprehension iterated chars. PEFT LoraConfig.target_modules accepts the bare-string convention. - ModelArchConfig gains `shared_expert_intermediate_size`. extract_arch_config reads `n_shared_experts` / `num_shared_experts` aliases and infers `n_shared_experts=1` when only `shared_expert_intermediate_size` is set. _compute_moe_mlp_elements and the structured + non-structured LoRA paths size the shared expert with its own intermediate (Qwen3.5-MoE: 512 vs routed moe_intermediate_size). - _determine_attention_impl_for_gpu_estimate uses copy.deepcopy so the resolver does not mutate nested text_config on the cached source. PreTrainedConfig._attn_implementation setter walks `sub_configs` and the prior shallow copy still touched the inner objects. * Studio: extend MoE/PLE/KV-share accounting to activation and skip-alias paths Five activation-path corrections plus two LoRA / skip-alias corrections so that shared-expert, per-layer-input, and KV-shared-layer support is symmetric across weights, LoRA, skip-quantizable, and activation paths. - _layer_qkv_mlp_sizes: include shared-expert FFN in mlp_size (live shared expert per token alongside routed experts) and keep K/V activation memory for KV-shared layers; only the WEIGHT path uses has_k/has_v from _layer_attention_dims. - _per_layer_activation_bytes / compute_activation_bytes: account for per_layer_input_gate (hd-sized) and per_layer_projection (pli-sized) per layer plus the global per_layer_model_projection [B,S,L,PLI] tensor when hidden_size_per_layer_input is set. - _build_text_module_elements: split mlp.experts into routed and mlp.shared_expert canonical entries; register layers.<i>.experts alias for Gemma4 enable_moe_block layouts and mlp.shared_experts (plural) alias for Exaone-MoE / Laguna / GLM4-MoE-Lite shared-expert variants. - _compute_moe_mlp_elements: split into _compute_routed_moe_elements and _compute_shared_moe_elements; only count shared_expert_gate (hd->1 Linear per shared expert) when shared_expert_intermediate_size is set, which is the Qwen2-MoE / Qwen3.5-MoE discriminator. Other shared-expert families (Exaone-MoE, HY-V3, GLM4-MoE-Lite, Laguna) lack the gate. - compute_lora_params: when target_modules='all-linear' bare keyword, drop routed and shared MoE expert LoRA contributions. PEFT's all-linear targets nn.Linear only; Unsloth's get_moe_target_parameters expands MoE expert nn.Parameter LoRA only when target_modules contains explicit gate_proj/up_proj/down_proj/gate_up_proj names. - _per_layer_input_lora_params: thread target_modules through and add the per-PLE-module contribution when the corresponding name appears, not only under all-linear. * Studio: top-k MoE activations, ERNIE list configs, suffix skips, multimodal full bytes Six estimator corrections aligning the detailed accounting paths with real training behavior: - _layer_qkv_mlp_sizes scales the MoE-layer mlp_size by num_experts_per_tok so the active routed-expert intermediate tensors are charged for activations. Adds num_experts_per_tok to ModelArchConfig and extracts it from num_experts_per_tok / top_k_experts (Gemma4 alias) in extract_arch_config. - compute_lora_params splits routed and shared MoE LoRA contributions so that bare target_modules='all-linear' zeroes routed (nn.Parameter expert tensors, which Unsloth's get_moe_target_parameters does NOT enable for the bare keyword) but keeps shared-expert LoRA (regular nn.Linear MLPs that Unsloth's get_peft_regex DOES match). - extract_arch_config gains a _first_scalar helper for ERNIE-style moe_intermediate_size = [routed, shared] lists, plus moe_num_experts and moe_num_shared_experts attribute aliases. When moe_intermediate_size is a pair and shared_expert_intermediate_size is unset, the second element is treated as the shared-expert intermediate. - estimate_required_model_memory_gb's detailed branch retains max(0, model_size_bytes - compute_total_params(arch) * 2) on top of the arch-derived breakdown.model_weights so multimodal models (vision/audio towers) and partially-modeled families (Gemma3n AltUp/Laurel etc.) do not silently drop bytes that the safetensors total includes. - _module_path_matches accepts a tail-only match when the skip entry is shorter than the alias path. Transformers' BNB quantizer suffix-matches short skip entries like ['q_proj'] / ['lm_head'] against full module paths; the previous len(skip) < len(alias) early-return missed those. - _per_layer_input_lora_params drops the all_linear branch and only counts PLE LoRA when the user explicitly names per_layer_input_gate / per_layer_projection / per_layer_model_projection. Unsloth's get_peft_regex requires module names to contain a component tag (mlp/attn/...); PLE module names lack any tag, so all-linear training does not attach LoRA to them. * Studio: full-FT extra optimizer/gradient inflation, MoE top-k aliases, ERNIE position dispatch, sibling experts aggregate When the safetensors total exceeds the text-arch fp16 estimate (multimodal vision/audio towers, partially-modeled families), only inflate the model weights line for adapter methods but extend optimizer + gradient bytes under full fine-tuning, where the extra params are trainable. DBRX exposes top-k routing as moe_top_k and Hunyuan-V1-MoE as moe_topk; neither is aliased to num_experts_per_tok via attribute_map, so probe both when extracting arch config. ERNIE 4.5 MoE / VL MoE configs declare MoE layers via moe_layer_start_index / moe_layer_end_index / moe_layer_interval (with -1 meaning the last layer); add the position-style dispatch alongside the existing mlp_layer_types / first_k_dense_replace / decoder_sparse_step paths. When moe_has_dense_mlp is set (Gemma4 enable_moe_block) the routed experts live as a sibling of self.mlp at layers.<i>.experts in the actual model layout; keep the layer mlp aggregate to the dense path and add a separate experts aggregate so a skip module model.layers.<i>.mlp does not collapse the routed experts as well. * Studio: extend MoE family extraction (Llama4 / DBRX / Hunyuan / ERNIE) and align dense vs routed MLP widths - Llama4: pick up `config.moe_layers` (auto-populated from interleave_moe_layer_step) so dense layer indices reflect the actual is_moe_layer dispatch. - Llama4: add a separate `dense_intermediate_size` derived from `intermediate_size_mlp` (used for the dense feed_forward path) and keep `intermediate_size` for the routed/shared expert width. Auto-attach one shared expert per MoE layer when the dense-vs-MoE width split is present. - DBRX: walk the `ffn_config` sub-config when extracting MoE attrs (moe_num_experts / moe_top_k / ffn_hidden_size). Without this DBRX is misclassified as a dense arch. - Hunyuan: normalize layer-wise `moe_topk` (and the canonical `num_experts_per_tok` lookup it shadows via attribute_map) through a worst-case scalar so the int(...) cast cannot crash on list values. - ERNIE 4.5 MoE: switch the start/end/interval dispatch to the model's `(layer_idx + 1) % interval == 0` modulo gate so MoE layers match the decoder when interval > 1. - ERNIE 4.5 VL MoE: drop the heuristic that read `moe_intermediate_size[1]` as the shared expert width; in VL configs [1] is the vision-routed width and shared experts are sized from [0]. - estimate_fp16_model_size_bytes: prefer the larger of config-derived and local-weight bytes so the multimodal extra_bytes correction can fire for local VLM directories. * Add tests for VRAM estimator extensions * Studio: trim verbose comments in VRAM estimator Collapse multi-paragraph rationale blocks to 1-3 lines stating the single load-bearing fact. Fix one inverted "fall through ... last" comment whose claim disagreed with the surrounding code. * Consolidate added tests into existing test_vram_estimation.py and test_gpu_selection.py Move Llama4 / DBRX / ERNIE arch-extraction tests into test_vram_estimation.py as TestLlama4ArchExtraction / TestDbrxFfnConfigExtraction / TestErniePhaseModuloDispatch / TestErnieVlSharedExpertWidth classes. Move estimate_fp16_model_size_bytes prefer-larger-of-config-or-local tests into test_gpu_selection.py as TestEstimateFp16ModelSizeBytesPrefersLocalWeights. Drop one redundant Llama4 num_dense_layers assertion already covered by the moe_layers dispatch test. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com> Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
186 lines
6.2 KiB
Markdown
186 lines
6.2 KiB
Markdown
# VRAM Estimation for Training
|
||
|
||
```
|
||
Total VRAM = Weights + LoRA Adapters + Optimizer + Gradients + Activations + CUDA Overhead
|
||
```
|
||
|
||
| Symbol | Meaning |
|
||
|--------|---------|
|
||
| `H` | `hidden_size` |
|
||
| `L` | `num_hidden_layers` |
|
||
| `V` | `vocab_size` |
|
||
| `K` | `(H / num_attention_heads) * num_key_value_heads` |
|
||
| `M` | `intermediate_size` (or `moe_intermediate_size`) |
|
||
| `E` | `num_experts` (1 for dense) |
|
||
| `r` | LoRA rank |
|
||
| `B` | `per_device_train_batch_size` |
|
||
| `S` | `max_seq_length` |
|
||
|
||
---
|
||
|
||
## 1. Model Weights
|
||
|
||
```
|
||
QKVO = (H + K + K + H) * H
|
||
MLP = H * M * 3 * E + (E * H if E > 1 else 0)
|
||
|
||
Quantizable = (QKVO + MLP) * L
|
||
Non-quantizable = 2*H*L + V*H + (V*H if not tie_embeddings else 0)
|
||
```
|
||
|
||
| Mode | Bytes |
|
||
|------|-------|
|
||
| QLoRA 4-bit | `Quantizable * 2 / 3.2 + Non-quantizable * 2` |
|
||
| LoRA / Full fp16 | `(Quantizable + Non-quantizable) * 2` |
|
||
|
||
The 3.2 factor (`16/5`) accounts for BNB NF4 blockwise scales. Repos whose
|
||
quantization config enables `bnb_4bit_use_double_quant` use a tighter, still
|
||
conservative 3.6 factor for the quantized portion of the weights.
|
||
When a 4-bit config has `llm_int8_skip_modules` entries that point to language
|
||
model layers or submodules, those quantizable weights are charged at fp16
|
||
instead of NF4. Generic embedding and multimodal skip names are already covered
|
||
by non-quantizable terms or excluded from text training weights.
|
||
|
||
## 2. LoRA Adapters
|
||
|
||
| Module | A | B |
|
||
|--------|---|---|
|
||
| q_proj | `H×r` | `r×H` |
|
||
| k_proj | `H×r` | `r×K` |
|
||
| v_proj | `H×r` | `r×K` |
|
||
| o_proj | `H×r` | `r×H` |
|
||
| gate_proj | `H×r` | `r×M` |
|
||
| up_proj | `H×r` | `r×M` |
|
||
| down_proj | `M×r` | `r×H` |
|
||
|
||
MLP modules multiply by `E` for MoE.
|
||
|
||
```
|
||
LoRA_bytes = sum(A + B per selected module) * L * 2
|
||
```
|
||
|
||
`all-linear` is treated as all known text linear modules in the table above.
|
||
The estimator deliberately does not infer multimodal or vision-tower LoRA
|
||
modules from config shapes; those modules vary too much across VLM families for
|
||
a generic config formula.
|
||
|
||
Some decoder configs expose layer-shape fields such as `layer_types`,
|
||
`head_dim`, `global_head_dim`, `num_global_key_value_heads`, `attention_k_eq_v`,
|
||
`num_kv_shared_layers`, `use_double_wide_mlp`, `vocab_size_per_layer_input`, and
|
||
`hidden_size_per_layer_input`. When those fields are present, the estimator
|
||
derives text weight and LoRA counts from the per-layer shapes instead of
|
||
assuming every layer has the same seven projection modules.
|
||
|
||
## 3. Optimizer States (calibrated)
|
||
|
||
| Optimizer | Bytes/param | Notes |
|
||
|-----------|------------|-------|
|
||
| `adamw_8bit` | 4 | BNB upcasts to fp32 during step |
|
||
| `adamw_torch` | 6 | Fused, no master copy |
|
||
| `paged_adamw_32bit` | 8 | Full fp32 states |
|
||
| `sgd` | 4 | |
|
||
|
||
Trainable params = all params (Full FT) or LoRA params only.
|
||
|
||
## 4. Gradients
|
||
|
||
```
|
||
Gradient_bytes = trainable_params * 2 (fp16, accumulated in-place)
|
||
```
|
||
|
||
## 5. Activations
|
||
|
||
Per-layer (from `unsloth_zoo/vllm_utils.py`):
|
||
```
|
||
Per_layer = (S*B*(H+K+K) + S*B*2 + S*B*(M+M)) * 2 * 1.25
|
||
```
|
||
|
||
When the resolved attention implementation is none of `flash_attention_2`,
|
||
`sdpa`, or `flex_attention` (PyTorch SDPA dispatches to flash or
|
||
memory-efficient kernels and FlexAttention is also a memory-efficient
|
||
kernel, all of which are O(n) in memory), activation memory also includes
|
||
a quadratic attention-score/workspace estimate:
|
||
|
||
```
|
||
Non_flash_attention = B * num_attention_heads * S^2 * 2 * 12.0 * effective_layers
|
||
Activations = max(Per_layer_with_gc, Non_flash_attention)
|
||
```
|
||
|
||
Studio resolves the attention implementation with Unsloth's
|
||
`resolve_attention_implementation` helper and uses that result directly. The
|
||
estimator does not duplicate model-family attention policy.
|
||
|
||
| GC Mode | Full FT | LoRA/QLoRA |
|
||
|---------|---------|------------|
|
||
| none | `L` layers | `L` layers |
|
||
| true (HF) | 2.0 | 1.0 |
|
||
| unsloth | 1.5 | 1.0 |
|
||
|
||
## 6. Floors
|
||
|
||
Activations use the computed formula directly:
|
||
|
||
```
|
||
activation_bytes = computed_activation_bytes
|
||
```
|
||
|
||
Full fine-tuning keeps the gradient floor at **15% of model weight memory** to
|
||
account for autograd overhead, NCCL buffers, mixed-precision scaling, and
|
||
PyTorch fragmentation:
|
||
|
||
```
|
||
gradient_bytes = max(computed_gradient_bytes, weights * 0.15)
|
||
```
|
||
|
||
For LoRA/QLoRA, the base model is frozen, so the weight-derived gradient floor
|
||
is capped by trainable-state and live-activation scale:
|
||
|
||
```
|
||
raw_gradient_bytes = trainable_params * 2
|
||
gradient_floor = min(weights * 0.15, max(computed_activation_bytes, optimizer_bytes))
|
||
gradient_bytes = max(raw_gradient_bytes, gradient_floor)
|
||
```
|
||
|
||
This prevents frozen quantized model size from dominating gradient/state
|
||
overhead when the measured runtime footprint is governed by LoRA optimizer
|
||
states and live activations.
|
||
|
||
## 7. CUDA Overhead
|
||
|
||
**1.4 GB** fixed — CUDA driver + PyTorch runtime, calibrated on RTX 5070 Ti.
|
||
|
||
## 8. Multi-GPU Overhead
|
||
|
||
When sharding across multiple GPUs, each additional GPU (beyond the first) contributes only **85%** of its free VRAM to the usable pool. The 15% discount accounts for NCCL all-reduce buffers, PCIe/NVLink transfer overhead, synchronization barriers, and memory fragmentation from non-uniform shard sizes. Calibrated empirically on 2-8 GPU setups with NVLink and PCIe topologies.
|
||
|
||
```
|
||
usable_gb = free[gpu_0] + sum(free[gpu_i] * 0.85 for i in 1..N)
|
||
```
|
||
|
||
---
|
||
|
||
## Parameter Flow
|
||
|
||
```
|
||
Frontend -> routes/{training,inference}.py
|
||
-> prepare_gpu_selection(gpu_ids, model_name, ...)
|
||
|
|
||
+-- gpu_ids is explicit (e.g. [5,6,7])
|
||
| -> resolve_requested_gpu_ids: validate against parent-visible set
|
||
| -> return all requested GPUs (model sharded across all of them)
|
||
|
|
||
+-- gpu_ids is None or []
|
||
-> auto_select_gpu_ids: estimate VRAM, pick minimum GPUs needed
|
||
-> estimate_required_model_memory_gb -> estimate_training_vram
|
||
-> greedy selection: rank GPUs by free VRAM, add until model fits
|
||
|
||
-> get_device_map(resolved_gpu_ids)
|
||
-> "balanced" if >1 GPU, "sequential" otherwise
|
||
|
||
-> worker subprocess: apply_gpu_ids(resolved_gpu_ids)
|
||
-> sets CUDA_VISIBLE_DEVICES before torch/CUDA init
|
||
```
|
||
|
||
Threaded params: `batch_size`, `max_seq_length`, `lora_r`, `target_modules`, `gradient_checkpointing`, `optim`.
|
||
|
||
Source: `studio/backend/utils/hardware/vram_estimation.py`
|