ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update VRAM estimator to cater to broader model configs (#5175)

Browse source 
* 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>
This commit is contained in:
Datta Nimmaturi 2026-05-05 16:42:36 +05:30 • committed by GitHub
commit 09505fcc6e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 2979 additions and 106 deletions
Download patch file Download diff file Expand all files Collapse all files

296
studio/backend/tests/test_gpu_selection.py
View file

@ -1049,6 +1049,182 @@ class TestMinGpuVram(unittest.TestCase):
class TestPerGpuFitGuardAllCounts(unittest.TestCase):
def test_training_estimate_resolves_attention_without_raising(self):
with (
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
patch(
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
return_value = (8 * (1024**3), "config"),
),
patch(
"utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
return_value = "unsloth/test",
),
patch(
"utils.hardware.hardware._load_config_for_gpu_estimate",
return_value = SimpleNamespace(
hidden_size = 4096,
num_hidden_layers = 32,
num_attention_heads = 32,
num_key_value_heads = 8,
intermediate_size = 14336,
vocab_size = 128256,
tie_word_embeddings = False,
),
),
patch(
"utils.hardware.hardware._determine_attention_impl_for_gpu_estimate",
return_value = "eager",
),
patch("utils.hardware.hardware.get_visible_gpu_count", return_value = 1),
):
_, metadata = estimate_required_model_memory_gb(
"unsloth/test",
training_type = "LoRA/QLoRA",
load_in_4bit = True,
)
self.assertEqual(metadata.get("estimation_mode"), "detailed")
self.assertEqual(metadata.get("attention_implementation"), "eager")
def test_training_estimate_falls_back_when_attention_resolution_fails(self):
with (
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
patch(
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
return_value = (8 * (1024**3), "config"),
),
patch(
"utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
return_value = "unsloth/test",
),
patch(
"utils.hardware.hardware._load_config_for_gpu_estimate",
return_value = SimpleNamespace(
hidden_size = 4096,
num_hidden_layers = 32,
num_attention_heads = 32,
num_key_value_heads = 8,
intermediate_size = 14336,
vocab_size = 128256,
tie_word_embeddings = False,
),
),
patch(
"utils.hardware.hardware._determine_attention_impl_for_gpu_estimate",
side_effect = RuntimeError("attention unavailable"),
),
patch("utils.hardware.hardware.get_visible_gpu_count", return_value = 1),
):
_, metadata = estimate_required_model_memory_gb(
"unsloth/test",
training_type = "LoRA/QLoRA",
load_in_4bit = True,
)
self.assertEqual(metadata.get("estimation_mode"), "detailed")
self.assertEqual(
metadata.get("attention_implementation"),
"eager",
)
def test_attention_resolver_does_not_mutate_loaded_config(self):
from utils.hardware import hardware as hardware_module
config = SimpleNamespace(
hidden_size = 1024,
num_hidden_layers = 2,
num_attention_heads = 8,
num_key_value_heads = 8,
intermediate_size = 2048,
vocab_size = 1024,
tie_word_embeddings = True,
)
def _stub_resolver(model_class, cfg):
cfg._attn_implementation = "eager"
return "eager"
with patch(
"unsloth.models._utils.resolve_attention_implementation",
side_effect = _stub_resolver,
):
hardware_module._determine_attention_impl_for_gpu_estimate(config)
self.assertFalse(hasattr(config, "_attn_implementation"))
def test_attention_resolver_handles_missing_model_mapping(self):
from utils.hardware import hardware as hardware_module
config = SimpleNamespace(
hidden_size = 1024,
num_hidden_layers = 2,
num_attention_heads = 8,
num_key_value_heads = 8,
intermediate_size = 2048,
vocab_size = 1024,
tie_word_embeddings = True,
)
captured = {}
def _stub_resolver(model_class, cfg):
captured["model_class"] = model_class
return "eager"
from transformers import AutoModel, AutoModelForCausalLM
with (
patch.object(AutoModelForCausalLM, "_model_mapping", new = None),
patch.object(AutoModel, "_model_mapping", new = None),
patch(
"unsloth.models._utils.resolve_attention_implementation",
side_effect = _stub_resolver,
),
):
result = hardware_module._determine_attention_impl_for_gpu_estimate(config)
self.assertEqual(result, "eager")
self.assertIsNone(captured["model_class"])
def test_attention_resolver_does_not_mutate_nested_text_config(self):
from utils.hardware import hardware as hardware_module
text_config = SimpleNamespace(
hidden_size = 1024,
num_hidden_layers = 2,
num_attention_heads = 8,
num_key_value_heads = 8,
intermediate_size = 2048,
vocab_size = 1024,
tie_word_embeddings = True,
)
config = SimpleNamespace(
hidden_size = 1024,
num_hidden_layers = 2,
num_attention_heads = 8,
num_key_value_heads = 8,
intermediate_size = 2048,
vocab_size = 1024,
tie_word_embeddings = True,
text_config = text_config,
)
def _stub_resolver(model_class, cfg):
cfg._attn_implementation = "eager"
inner = getattr(cfg, "text_config", None)
if inner is not None:
inner._attn_implementation = "eager"
return "eager"
with patch(
"unsloth.models._utils.resolve_attention_implementation",
side_effect = _stub_resolver,
):
hardware_module._determine_attention_impl_for_gpu_estimate(config)
self.assertFalse(hasattr(config, "_attn_implementation"))
self.assertFalse(hasattr(text_config, "_attn_implementation"))
def test_min_per_gpu_generated_for_all_visible_counts(self):
with (
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
@ -1125,3 +1301,123 @@ class TestXpuRejection(_GpuCacheResetMixin, unittest.TestCase):
with patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU):
with self.assertRaisesRegex(ValueError, "only supported on CUDA"):
prepare_gpu_selection([0], model_name = "unsloth/test")
class TestEstimateFp16ModelSizeBytesPrefersLocalWeights(unittest.TestCase):
def _run(
self,
model_path,
*,
config_bytes,
local_bytes,
safetensors_params = None,
config = object(),
):
from utils.hardware import hardware as hardware_module
with (
patch.object(
hardware_module,
"_resolve_model_identifier_for_gpu_estimate",
return_value = model_path,
),
patch.object(
hardware_module,
"_get_hf_safetensors_total_params",
return_value = safetensors_params,
),
patch.object(
hardware_module,
"_load_config_for_gpu_estimate",
return_value = config,
),
patch.object(
hardware_module,
"_estimate_fp16_model_size_bytes_from_config",
return_value = config_bytes,
),
patch.object(
hardware_module,
"_get_local_weight_size_bytes",
return_value = local_bytes,
),
):
return hardware_module.estimate_fp16_model_size_bytes(model_path)
def test_local_weight_bytes_preferred_when_larger_than_config(self):
bytes_, src = self._run(
"/local/vlm",
config_bytes = 2 * (1 << 30),
local_bytes = 20 * (1 << 30),
)
self.assertEqual(bytes_, 20 * (1 << 30))
self.assertEqual(src, "weight_bytes")
def test_config_bytes_preferred_when_larger_than_local(self):
bytes_, src = self._run(
"/local/text-only",
config_bytes = 20 * (1 << 30),
local_bytes = 2 * (1 << 30),
)
self.assertEqual(bytes_, 20 * (1 << 30))
self.assertEqual(src, "config")
def test_config_bytes_returned_when_no_local_weights(self):
bytes_, src = self._run(
"/local/no-weights",
config_bytes = 5 * (1 << 30),
local_bytes = None,
)
self.assertEqual(bytes_, 5 * (1 << 30))
self.assertEqual(src, "config")
def test_local_bytes_returned_when_config_resolution_fails(self):
bytes_, src = self._run(
"/local/no-config",
config_bytes = None,
local_bytes = 7 * (1 << 30),
config = None,
)
self.assertEqual(bytes_, 7 * (1 << 30))
self.assertEqual(src, "weight_bytes")
def test_equal_local_and_config_keeps_config_label(self):
# why: tie-breaker is "local must be strictly larger" so an exact
# match keeps the config-derived path.
same = 8 * (1 << 30)
bytes_, src = self._run(
"/local/equal",
config_bytes = same,
local_bytes = same,
)
self.assertEqual(bytes_, same)
self.assertEqual(src, "config")
def test_remote_safetensors_path_unaffected_by_local_weights(self):
from utils.hardware import hardware as hardware_module
with (
patch.object(
hardware_module,
"_resolve_model_identifier_for_gpu_estimate",
return_value = "owner/repo",
),
patch.object(
hardware_module,
"_get_hf_safetensors_total_params",
return_value = 1_000_000_000,
),
patch.object(
hardware_module,
"_load_config_for_gpu_estimate",
) as mock_load,
patch.object(
hardware_module,
"_get_local_weight_size_bytes",
) as mock_local,
):
bytes_, src = hardware_module.estimate_fp16_model_size_bytes("owner/repo")
self.assertEqual(bytes_, 2 * 1_000_000_000)
self.assertEqual(src, "safetensors")
mock_load.assert_not_called()
mock_local.assert_not_called()

1674
studio/backend/tests/test_vram_estimation.py
View file

File diff suppressed because it is too large Load diff

89
studio/backend/utils/hardware/VRAM_ESTIMATION.md
View file

@ -33,7 +33,13 @@ Non-quantizable = 2*H*L + V*H + (V*H if not tie_embeddings else 0)
| 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.
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
@ -53,6 +59,18 @@ 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 |
@ -77,6 +95,21 @@ 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 |
@ -85,13 +118,33 @@ Per_layer = (S*B*(H+K+K) + S*B*2 + S*B*(M+M)) * 2 * 1.25
## 6. Floors
Gradients and activations have minimum floors at **15% of model weight memory** to account for autograd overhead, attention score matrices, NCCL buffers, mixed-precision scaling, and PyTorch fragmentation.
Activations use the computed formula directly:
```
gradient_bytes = max(computed, weights * 0.15)
activation_bytes = max(computed, weights * 0.15 * B/2)
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.
@ -106,34 +159,6 @@ usable_gb = free[gpu_0] + sum(free[gpu_i] * 0.85 for i in 1..N)
---
## Reference Table (bsz=2, seq=2048, rank=16, GC=unsloth, adamw_8bit)
| Model | Weights | LoRA | Optim | Grad | Act | CUDA | Total |
|-------|---------|------|-------|------|-----|------|-------|
| 0.5B QLoRA | 0.5 | 0.0 | 0.0 | 0.1 | 0.1 | 1.4 | **2.1** |
| 1B QLoRA | 1.1 | 0.0 | 0.0 | 0.2 | 0.2 | 1.4 | **2.9** |
| 3B QLoRA | 2.4 | 0.0 | 0.1 | 0.5 | 0.5 | 1.4 | **4.9** |
| 8B QLoRA | 6.0 | 0.1 | 0.2 | 1.2 | 1.2 | 1.4 | **10.1** |
| 8B LoRA fp16 | 15.0 | 0.1 | 0.2 | 3.0 | 3.0 | 1.4 | **22.6** |
| 8B Full FT | 15.0 | — | 29.9 | 15.0 | 3.0 | 1.4 | **64.2** |
| 32B LoRA fp16 | 61.0 | 0.2 | 0.5 | 12.2 | 12.2 | 1.4 | **87.6** |
| 72B QLoRA | 45.5 | 0.4 | 0.8 | 9.1 | 9.1 | 1.4 | **66.3** |
## E2E Validation (Llama-3.2-1B, B200 emulating 24GB)
| Config | Estimated | Actual (nvsmi) | Error |
|--------|----------|----------------|-------|
| QLoRA bsz=2 seq=512 | 2.55 GB | 2.65 GB | -3.7% |
| QLoRA bsz=2 seq=2048 | 2.60 GB | 2.65 GB | -1.8% |
| QLoRA bsz=4 seq=2048 | 2.65 GB | 2.65 GB | +0.0% |
| LoRA fp16 bsz=2 | 3.84 GB | 3.88 GB | -1.0% |
| Full FT adamw_8bit | 10.89 GB | 10.80 GB | +0.8% |
| Full FT adamw_torch | 13.19 GB | 12.93 GB | +2.0% |
*Note: e2e numbers predate the 15% floors, which add safety margin on top.*
---
## Parameter Flow
```

75
studio/backend/utils/hardware/hardware.py
View file

@ -774,6 +774,34 @@ def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = Non
return None
def _determine_attention_impl_for_gpu_estimate(config) -> str:
import copy as _copy
from unsloth.models._utils import resolve_attention_implementation
from transformers import AutoModel, AutoModelForCausalLM
# why: resolve_attention_implementation calls _set_attn_impl which writes
# _attn_implementation onto the config; PreTrainedConfig's setter walks
# `sub_configs` and propagates to nested text_config / sub-configs, so a
# shallow copy still mutates those shared inner objects on the cached
# config returned by _load_config_for_gpu_estimate. Deepcopy isolates them.
config_copy = _copy.deepcopy(config)
model_class = None
for auto_model in (AutoModelForCausalLM, AutoModel):
mapping = getattr(auto_model, "_model_mapping", None)
if mapping is None:
continue
try:
if config_copy.__class__ in mapping:
model_class = mapping[config_copy.__class__]
break
except Exception:
continue
return resolve_attention_implementation(model_class, config_copy)
def _estimate_fp16_model_size_bytes_from_config(config) -> Optional[int]:
from .vram_estimation import extract_arch_config, compute_total_params
@ -844,12 +872,21 @@ def estimate_fp16_model_size_bytes(
return int(total_params * 2), "safetensors"
config = _load_config_for_gpu_estimate(estimate_model, hf_token = hf_token)
config_bytes: Optional[int] = None
if config is not None:
config_bytes = _estimate_fp16_model_size_bytes_from_config(config)
if config_bytes is not None:
return config_bytes, "config"
local_bytes = _get_local_weight_size_bytes(estimate_model)
# why: config-derived bytes cover only the text tower; local safetensors
# include vision/audio towers. Take the larger so the multimodal
# extra_bytes correction can fire.
if config_bytes is not None and local_bytes is not None:
if local_bytes > config_bytes:
return local_bytes, "weight_bytes"
return config_bytes, "config"
if config_bytes is not None:
return config_bytes, "config"
if local_bytes is not None:
return local_bytes, "weight_bytes"
@ -877,6 +914,9 @@ def estimate_required_model_memory_gb(
TrainingVramConfig,
extract_arch_config,
estimate_training_vram,
compute_total_params,
compute_optimizer_bytes,
compute_gradient_bytes,
CUDA_OVERHEAD_BYTES,
QUANT_4BIT_FACTOR,
DEFAULT_TARGET_MODULES,
@ -926,13 +966,44 @@ def estimate_required_model_memory_gb(
model_name, hf_token = hf_token
)
config = _load_config_for_gpu_estimate(estimate_model, hf_token = hf_token)
if config is not None:
try:
vram_config.attention_implementation = (
_determine_attention_impl_for_gpu_estimate(config)
)
except Exception as e:
logger.warning(
"Could not resolve attention implementation for '%s': %s",
estimate_model,
e,
)
# why: if we cannot prove flash attention is usable, charge the
# quadratic non-flash activation path so GPU selection stays
# conservative.
vram_config.attention_implementation = "eager"
arch = extract_arch_config(config) if config is not None else None
if arch is not None:
breakdown = estimate_training_vram(arch, vram_config)
# why: extract_arch_config only sees text_config; safetensors include
# vision/audio tower bytes that the text-arch fp16 total misses.
arch_fp16_bytes = compute_total_params(arch) * 2
extra_bytes = max(0, int(model_size_bytes) - arch_fp16_bytes)
if extra_bytes > 0:
breakdown.model_weights += extra_bytes
if training_method == "full":
# why: full fine-tuning makes the extra (vision/audio) params
# trainable; optimizer + gradient bytes scale with them too.
extra_params = extra_bytes // 2
breakdown.optimizer_states += compute_optimizer_bytes(
extra_params,
vram_config.optimizer,
)
breakdown.gradients += compute_gradient_bytes(extra_params)
required_gb = breakdown.total / (1024**3)
metadata["required_gb"] = round(required_gb, 3)
metadata["estimation_mode"] = "detailed"
metadata["attention_implementation"] = vram_config.attention_implementation
metadata["vram_breakdown"] = breakdown.to_gb_dict()
max_gpus = max(1, get_visible_gpu_count())
for n_gpus in range(1, max_gpus + 1):

955
studio/backend/utils/hardware/vram_estimation.py
View file

File diff suppressed because it is too large Load diff