diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index fb4f38565b..a1fe5653ef 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -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() diff --git a/studio/backend/tests/test_vram_estimation.py b/studio/backend/tests/test_vram_estimation.py index 0be067310d..e54ae6dcf8 100644 --- a/studio/backend/tests/test_vram_estimation.py +++ b/studio/backend/tests/test_vram_estimation.py @@ -2,7 +2,9 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. import unittest +from dataclasses import replace from types import SimpleNamespace +from unittest.mock import patch from utils.hardware.vram_estimation import ( ModelArchConfig, @@ -116,6 +118,55 @@ GPT_OSS = ModelArchConfig( num_dense_layers = 0, ) +STRUCTURED_MIXED = ModelArchConfig( + hidden_size = 256, + num_hidden_layers = 6, + num_attention_heads = 4, + num_key_value_heads = 2, + intermediate_size = 512, + vocab_size = 1024, + tie_word_embeddings = True, + head_dim = 80, + global_head_dim = 96, + num_global_key_value_heads = 1, + attention_k_eq_v = True, + layer_types = [ + "sliding_attention", + "full_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + "full_attention", + ], +) + +STRUCTURED_SHARED = ModelArchConfig( + hidden_size = 192, + num_hidden_layers = 4, + num_attention_heads = 6, + num_key_value_heads = 2, + intermediate_size = 384, + vocab_size = 512, + tie_word_embeddings = True, + head_dim = 32, + num_kv_shared_layers = 2, + use_double_wide_mlp = True, + vocab_size_per_layer_input = 128, + hidden_size_per_layer_input = 48, + quant_4bit_factor = 3.6, +) + +QUANT_SKIP_STRUCTURED = replace( + STRUCTURED_SHARED, + quantization_skip_modules = [ + "model.layers.0.self_attn.q_proj", + "language_model.model.layers.1.mlp", + "layers.2", + "vision_tower", + "embed_tokens", + ], +) + class TestExtractArchConfig(unittest.TestCase): def test_basic_config(self): @@ -182,6 +233,42 @@ class TestExtractArchConfig(unittest.TestCase): arch = extract_arch_config(hf_config) self.assertEqual(arch.intermediate_size, 8192) + def test_structural_and_quantization_fields_are_config_derived(self): + hf_config = SimpleNamespace( + hidden_size = 256, + num_hidden_layers = 2, + num_attention_heads = 4, + num_key_value_heads = 2, + intermediate_size = 512, + vocab_size = 1024, + tie_word_embeddings = True, + head_dim = 80, + global_head_dim = 96, + num_global_key_value_heads = 1, + attention_k_eq_v = True, + layer_types = ["sliding_attention", "full_attention"], + num_kv_shared_layers = 1, + use_double_wide_mlp = True, + vocab_size_per_layer_input = 128, + hidden_size_per_layer_input = 48, + quantization_config = { + "bnb_4bit_use_double_quant": True, + "llm_int8_skip_modules": ["model.layers.0.self_attn"], + }, + ) + arch = extract_arch_config(hf_config) + self.assertEqual(arch.head_dim, 80) + self.assertEqual(arch.global_head_dim, 96) + self.assertEqual(arch.num_global_key_value_heads, 1) + self.assertTrue(arch.attention_k_eq_v) + self.assertEqual(arch.layer_types, ["sliding_attention", "full_attention"]) + self.assertEqual(arch.num_kv_shared_layers, 1) + self.assertTrue(arch.use_double_wide_mlp) + self.assertEqual(arch.vocab_size_per_layer_input, 128) + self.assertEqual(arch.hidden_size_per_layer_input, 48) + self.assertEqual(arch.quantization_skip_modules, ["model.layers.0.self_attn"]) + self.assertEqual(arch.quant_4bit_factor, 3.6) + class TestModelWeightsBytes(unittest.TestCase): def test_llama_8b_fp16(self): @@ -238,6 +325,18 @@ class TestLoraParams(unittest.TestCase): ratio = moe_lora / dense_lora self.assertAlmostEqual(ratio, 8.0, delta = 0.5) + def test_structured_moe_mlp_modules_scale_with_experts(self): + structured_moe = replace(QWEN3_MOE_30B, head_dim = 128) + dense_like = replace( + structured_moe, + num_experts = None, + moe_intermediate_size = None, + ) + target_modules = ["gate_proj", "up_proj", "down_proj"] + dense_lora = compute_lora_params(dense_like, 16, target_modules) + moe_lora = compute_lora_params(structured_moe, 16, target_modules) + self.assertGreater(moe_lora, dense_lora * 20) + def test_attention_modules_same_for_moe(self): dense_attn = compute_lora_params( LLAMA_8B, 16, ["q_proj", "k_proj", "v_proj", "o_proj"] @@ -247,6 +346,41 @@ class TestLoraParams(unittest.TestCase): ) self.assertEqual(dense_attn, moe_attn) + def test_all_linear_uses_default_text_modules(self): + text_only = compute_lora_params(STRUCTURED_MIXED, 16, DEFAULT_TARGET_MODULES) + all_linear = compute_lora_params(STRUCTURED_MIXED, 16, ["all-linear"]) + self.assertEqual(all_linear, text_only) + + def test_structural_layer_shapes_are_config_driven(self): + unstructured_arch = replace( + STRUCTURED_MIXED, + head_dim = None, + global_head_dim = None, + num_global_key_value_heads = None, + attention_k_eq_v = False, + layer_types = None, + ) + self.assertNotEqual( + compute_lora_params(unstructured_arch, 16, ["all-linear"]), + compute_lora_params(STRUCTURED_MIXED, 16, ["all-linear"]), + ) + self.assertNotEqual( + compute_model_weights_bytes(unstructured_arch, "qlora", True), + compute_model_weights_bytes(STRUCTURED_MIXED, "qlora", True), + ) + + def test_shared_kv_and_per_layer_inputs_change_weight_count(self): + unstructured_arch = replace( + STRUCTURED_SHARED, + head_dim = None, + num_kv_shared_layers = 0, + use_double_wide_mlp = False, + ) + self.assertNotEqual( + compute_model_weights_bytes(unstructured_arch, "qlora", True), + compute_model_weights_bytes(STRUCTURED_SHARED, "qlora", True), + ) + class TestOptimizerBytes(unittest.TestCase): def test_adamw_8bit(self): @@ -293,6 +427,163 @@ class TestActivationBytes(unittest.TestCase): act_4k = compute_activation_bytes(LLAMA_8B, 2, 4096, "unsloth") self.assertAlmostEqual(act_4k / act_2k, 2.0, delta = 0.1) + def test_flash_attention_uses_linear_path(self): + flash = compute_activation_bytes( + STRUCTURED_MIXED, + 1, + 4096, + "unsloth", + is_lora = True, + attention_implementation = "flash_attention_2", + ) + default = compute_activation_bytes( + STRUCTURED_MIXED, + 1, + 4096, + "unsloth", + is_lora = True, + ) + self.assertEqual(flash, default) + + def test_sdpa_attention_uses_linear_path(self): + flash = compute_activation_bytes( + STRUCTURED_MIXED, + 1, + 4096, + "unsloth", + is_lora = True, + attention_implementation = "flash_attention_2", + ) + sdpa = compute_activation_bytes( + STRUCTURED_MIXED, + 1, + 4096, + "unsloth", + is_lora = True, + attention_implementation = "sdpa", + ) + self.assertEqual(sdpa, flash) + + def test_non_flash_attention_uses_quadratic_path(self): + seq_len = 4096 + expected_quadratic = ( + 1 * STRUCTURED_MIXED.num_attention_heads * seq_len * seq_len * 2 * 12.0 + ) + for attention_implementation in ("eager", "unknown_impl", None): + with self.subTest(attention_implementation = attention_implementation): + non_flash = compute_activation_bytes( + STRUCTURED_MIXED, + 1, + seq_len, + "unsloth", + is_lora = True, + attention_implementation = attention_implementation, + ) + self.assertEqual(non_flash, int(expected_quadratic)) + + def test_non_flash_attention_without_gc_scales_quadratic_path_by_layers(self): + seq_len = 4096 + one_layer = ( + 1 * STRUCTURED_MIXED.num_attention_heads * seq_len * seq_len * 2 * 12.0 + ) + non_flash = compute_activation_bytes( + STRUCTURED_MIXED, + 1, + seq_len, + "none", + is_lora = True, + attention_implementation = "eager", + ) + self.assertEqual(non_flash, int(one_layer * STRUCTURED_MIXED.num_hidden_layers)) + self.assertGreater(non_flash, int(one_layer)) + + +class TestQuantizationSkips(unittest.TestCase): + def test_skipped_language_layers_stay_fp16(self): + no_skips = replace(QUANT_SKIP_STRUCTURED, quantization_skip_modules = []) + skipped = compute_model_weights_bytes(QUANT_SKIP_STRUCTURED, "qlora", True) + quantized = compute_model_weights_bytes(no_skips, "qlora", True) + self.assertGreater(skipped, quantized) + + def test_non_language_skips_do_not_double_count_text_weights(self): + arch = replace( + QUANT_SKIP_STRUCTURED, + quantization_skip_modules = ["vision_tower", "embed_tokens"], + ) + no_skips = replace(QUANT_SKIP_STRUCTURED, quantization_skip_modules = []) + self.assertEqual( + compute_model_weights_bytes(arch, "qlora", True), + compute_model_weights_bytes(no_skips, "qlora", True), + ) + + def test_double_quant_factor_reduces_quantized_weight_storage(self): + default_quant = replace(STRUCTURED_MIXED, quant_4bit_factor = 16 / 5) + double_quant = replace(STRUCTURED_MIXED, quant_4bit_factor = 3.6) + self.assertLess( + compute_model_weights_bytes(double_quant, "qlora", True), + compute_model_weights_bytes(default_quant, "qlora", True), + ) + + def test_prefixed_parent_and_child_skips_do_not_double_count(self): + parent_only = replace( + QUANT_SKIP_STRUCTURED, + quantization_skip_modules = ["language_model.model.layers.1.mlp"], + ) + parent_and_child = replace( + QUANT_SKIP_STRUCTURED, + quantization_skip_modules = [ + "language_model.model.layers.1.mlp", + "language_model.model.layers.1.mlp.gate_proj", + "model.layers.1.mlp.up_proj", + ], + ) + self.assertEqual( + compute_model_weights_bytes(parent_and_child, "qlora", True), + compute_model_weights_bytes(parent_only, "qlora", True), + ) + + def test_vlm_prefix_skip_module_does_not_match_text_alias(self): + # vision_tower-prefixed skips must not shadow text aliases sharing the + # same suffix. + baseline = replace(QUANT_SKIP_STRUCTURED, quantization_skip_modules = []) + vlm_skip = replace( + QUANT_SKIP_STRUCTURED, + quantization_skip_modules = [ + "vision_tower.model.layers.0.self_attn.q_proj", + "vision_tower.model.layers.1.mlp", + ], + ) + self.assertEqual( + compute_model_weights_bytes(vlm_skip, "qlora", True), + compute_model_weights_bytes(baseline, "qlora", True), + ) + + def test_mla_skip_module_uses_authoritative_attn_total(self): + from utils.hardware.vram_estimation import ( + _build_text_module_elements, + _compute_attn_elements, + ) + + mla = ModelArchConfig( + hidden_size = 2048, + num_hidden_layers = 4, + num_attention_heads = 16, + num_key_value_heads = 16, + intermediate_size = 8192, + vocab_size = 32000, + tie_word_embeddings = False, + q_lora_rank = 512, + kv_lora_rank = 128, + qk_nope_head_dim = 64, + qk_rope_head_dim = 32, + v_head_dim = 64, + ) + elements, _ = _build_text_module_elements(mla) + self.assertEqual( + elements["text.layers.0.self_attn"], + _compute_attn_elements(mla), + ) + class TestEstimateTrainingVram(unittest.TestCase): def test_llama_8b_qlora_reasonable_total(self): @@ -430,6 +721,90 @@ class TestEstimateTrainingVram(unittest.TestCase): v32.optimizer_states / v8.optimizer_states, 1.5, delta = 0.1 ) + def test_min_gpu_vram_treats_activations_as_per_gpu_fixed(self): + config = TrainingVramConfig(training_method = "qlora", load_in_4bit = True) + breakdown = estimate_training_vram(LLAMA_8B, config) + shardable = ( + breakdown.model_weights + + breakdown.lora_adapters + + breakdown.optimizer_states + + breakdown.gradients + ) + per_gpu_fixed = breakdown.activations + breakdown.cuda_overhead + for n_gpus in (1, 2, 4): + self.assertEqual( + breakdown.min_gpu_vram(n_gpus), + shardable // n_gpus + per_gpu_fixed, + ) + + def test_qlora_gradient_floor_is_capped_by_trainable_scale(self): + config = TrainingVramConfig( + training_method = "qlora", + batch_size = 1, + max_seq_length = 512, + lora_rank = 16, + target_modules = ["all-linear"], + gradient_checkpointing = "unsloth", + optimizer = "adamw_8bit", + load_in_4bit = True, + ) + breakdown = estimate_training_vram(LLAMA_8B, config) + lora_params = compute_lora_params(LLAMA_8B, 16, DEFAULT_TARGET_MODULES) + optimizer_bytes = compute_optimizer_bytes(lora_params, "adamw_8bit") + weight_floor = int(breakdown.model_weights * 0.15) + + self.assertEqual( + breakdown.gradients, + max(breakdown.activations_computed, optimizer_bytes), + ) + self.assertLess(breakdown.gradients, weight_floor) + self.assertEqual(breakdown.activations, breakdown.activations_computed) + + def test_full_finetuning_gradient_floor_remains_uncapped(self): + config = TrainingVramConfig( + training_method = "full", + batch_size = 1, + max_seq_length = 512, + gradient_checkpointing = "unsloth", + optimizer = "adamw_8bit", + load_in_4bit = False, + ) + expected_floor = int( + compute_model_weights_bytes(LLAMA_8B, "full", False) * 0.15 + ) + with patch( + "utils.hardware.vram_estimation.compute_gradient_bytes", + return_value = 1, + ): + breakdown = estimate_training_vram(LLAMA_8B, config) + self.assertEqual(breakdown.gradients, expected_floor) + + def test_non_flash_attention_flows_into_training_estimate(self): + config = TrainingVramConfig( + training_method = "qlora", + batch_size = 1, + max_seq_length = 4096, + lora_rank = 16, + target_modules = ["all-linear"], + gradient_checkpointing = "unsloth", + optimizer = "adamw_8bit", + load_in_4bit = True, + attention_implementation = "eager", + ) + breakdown = estimate_training_vram(STRUCTURED_MIXED, config) + self.assertEqual(breakdown.activations, breakdown.activations_computed) + self.assertGreater( + breakdown.activations, + compute_activation_bytes( + STRUCTURED_MIXED, + 1, + 4096, + "unsloth", + is_lora = True, + attention_implementation = "flash_attention_2", + ), + ) + class TestExtractArchConfigMoE(unittest.TestCase): def test_deepseek_v3_shared_experts(self): @@ -471,11 +846,16 @@ class TestExtractArchConfigMoE(unittest.TestCase): moe_intermediate_size = 768, decoder_sparse_step = 1, mlp_only_layers = [], + head_dim = 128, ) arch = extract_arch_config(hf_config) self.assertEqual(arch.num_experts, 128) self.assertEqual(arch.num_dense_layers, 0) + self.assertEqual(arch.head_dim, 128) self.assertIsNone(arch.q_lora_rank) + total_b = compute_total_params(arch) / 1e9 + self.assertGreater(total_b, 20) + self.assertLess(total_b, 50) def test_qwen3_moe_with_mlp_only_layers(self): hf_config = SimpleNamespace( @@ -542,6 +922,343 @@ class TestExtractArchConfigMoE(unittest.TestCase): self.assertEqual(arch.n_shared_experts, 0) self.assertEqual(arch.num_dense_layers, 0) self.assertIsNone(arch.q_lora_rank) + self.assertFalse(arch.moe_has_dense_mlp) + + def test_enable_moe_block_extracted_as_moe_has_dense_mlp(self): + hf_config = SimpleNamespace( + hidden_size = 2048, + num_hidden_layers = 8, + num_attention_heads = 16, + num_key_value_heads = 4, + intermediate_size = 4096, + vocab_size = 32000, + tie_word_embeddings = True, + num_experts = 8, + moe_intermediate_size = 1024, + head_dim = 128, + layer_types = ["full_attention"] * 8, + enable_moe_block = True, + ) + arch = extract_arch_config(hf_config) + self.assertTrue(arch.moe_has_dense_mlp) + + +class TestParallelDenseMoE(unittest.TestCase): + def _arch(self, **overrides): + base = ModelArchConfig( + hidden_size = 512, + num_hidden_layers = 4, + num_attention_heads = 8, + num_key_value_heads = 2, + intermediate_size = 1024, + vocab_size = 1024, + tie_word_embeddings = True, + num_experts = 8, + moe_intermediate_size = 512, + num_dense_layers = 0, + head_dim = 64, + layer_types = ["full_attention"] * 4, + ) + return replace(base, **overrides) + + def test_total_params_includes_parallel_dense_when_enable_moe_block(self): + without_parallel = self._arch(moe_has_dense_mlp = False) + with_parallel = self._arch(moe_has_dense_mlp = True) + self.assertGreater( + compute_total_params(with_parallel), + compute_total_params(without_parallel), + ) + + def test_lora_params_includes_parallel_dense_when_enable_moe_block(self): + without_parallel = self._arch(moe_has_dense_mlp = False) + with_parallel = self._arch(moe_has_dense_mlp = True) + target = ["gate_proj", "up_proj", "down_proj"] + self.assertGreater( + compute_lora_params(with_parallel, 16, target), + compute_lora_params(without_parallel, 16, target), + ) + + def test_activation_bytes_includes_parallel_dense_when_enable_moe_block(self): + without_parallel = self._arch(moe_has_dense_mlp = False) + with_parallel = self._arch(moe_has_dense_mlp = True) + self.assertGreater( + compute_activation_bytes( + with_parallel, + 1, + 2048, + "unsloth", + is_lora = True, + ), + compute_activation_bytes( + without_parallel, + 1, + 2048, + "unsloth", + is_lora = True, + ), + ) + + def test_layer_aggregates_split_dense_mlp_from_experts(self): + from utils.hardware.vram_estimation import _build_text_module_elements + + with_parallel = self._arch(moe_has_dense_mlp = True) + elements, _ = _build_text_module_elements(with_parallel) + moe_only = ( + with_parallel.hidden_size + * with_parallel.moe_intermediate_size + * 3 + * with_parallel.num_experts + + with_parallel.num_experts * with_parallel.hidden_size + ) + dense_only = with_parallel.hidden_size * with_parallel.intermediate_size * 3 + # why: under gemma4 enable_moe_block, the layer's `self.experts` is a + # sibling of `self.mlp`; the `text.layers..mlp` aggregate must + # cover the dense path only, with experts in their own aggregate. + self.assertEqual(elements["text.layers.0.mlp"], dense_only) + self.assertEqual(elements["text.layers.0.experts"], moe_only) + + +class TestDenseLayerIndices(unittest.TestCase): + def test_non_prefix_mlp_only_layers_preserve_position(self): + hf_config = SimpleNamespace( + hidden_size = 1024, + num_hidden_layers = 8, + num_attention_heads = 16, + num_key_value_heads = 4, + intermediate_size = 2048, + vocab_size = 32000, + tie_word_embeddings = True, + num_local_experts = 4, + moe_intermediate_size = 512, + decoder_sparse_step = 1, + mlp_only_layers = [3, 5], + ) + arch = extract_arch_config(hf_config) + self.assertEqual(arch.num_dense_layers, 2) + self.assertIn(3, arch.dense_layer_indices) + self.assertIn(5, arch.dense_layer_indices) + self.assertNotIn(0, arch.dense_layer_indices) + + def test_first_k_dense_replace_indices_are_prefix(self): + hf_config = SimpleNamespace( + hidden_size = 1024, + num_hidden_layers = 6, + num_attention_heads = 16, + num_key_value_heads = 4, + intermediate_size = 2048, + vocab_size = 32000, + tie_word_embeddings = False, + n_routed_experts = 8, + moe_intermediate_size = 512, + first_k_dense_replace = 2, + ) + arch = extract_arch_config(hf_config) + self.assertEqual(tuple(arch.dense_layer_indices), (0, 1)) + + +class TestKvSharedLayer(unittest.TestCase): + def test_fully_shared_kv_returns_false_matching_upstream(self): + from utils.hardware.vram_estimation import _is_kv_shared_layer + + arch = ModelArchConfig( + hidden_size = 512, + num_hidden_layers = 4, + num_attention_heads = 8, + num_key_value_heads = 2, + intermediate_size = 1024, + vocab_size = 1024, + num_kv_shared_layers = 4, + ) + for i in range(arch.num_hidden_layers): + self.assertFalse(_is_kv_shared_layer(arch, i)) + + def test_partial_share_returns_true_for_tail_layers(self): + from utils.hardware.vram_estimation import _is_kv_shared_layer + + arch = ModelArchConfig( + hidden_size = 512, + num_hidden_layers = 4, + num_attention_heads = 8, + num_key_value_heads = 2, + intermediate_size = 1024, + vocab_size = 1024, + num_kv_shared_layers = 2, + ) + self.assertFalse(_is_kv_shared_layer(arch, 0)) + self.assertFalse(_is_kv_shared_layer(arch, 1)) + self.assertTrue(_is_kv_shared_layer(arch, 2)) + self.assertTrue(_is_kv_shared_layer(arch, 3)) + + +class TestFlexAttentionLinear(unittest.TestCase): + def test_flex_attention_treated_as_linear(self): + flash = compute_activation_bytes( + STRUCTURED_MIXED, + 1, + 4096, + "unsloth", + is_lora = True, + attention_implementation = "flash_attention_2", + ) + flex = compute_activation_bytes( + STRUCTURED_MIXED, + 1, + 4096, + "unsloth", + is_lora = True, + attention_implementation = "flex_attention", + ) + self.assertEqual(flex, flash) + + +class TestNonStructuredParallelDense(unittest.TestCase): + def _arch(self, **overrides): + base = ModelArchConfig( + hidden_size = 1024, + num_hidden_layers = 4, + num_attention_heads = 16, + num_key_value_heads = 4, + intermediate_size = 4096, + vocab_size = 32000, + tie_word_embeddings = False, + num_experts = 8, + moe_intermediate_size = 768, + num_dense_layers = 0, + moe_has_dense_mlp = True, + ) + return replace(base, **overrides) + + def test_skip_module_uses_intermediate_size_for_parallel_dense(self): + from utils.hardware.vram_estimation import _build_text_module_elements + + arch = self._arch() + elements, _ = _build_text_module_elements(arch) + gate_proj = elements["text.layers.0.mlp.gate_proj"] + self.assertEqual(gate_proj, arch.hidden_size * arch.intermediate_size) + + +class TestPerLayerInputAccounting(unittest.TestCase): + def _arch(self, **overrides): + base = ModelArchConfig( + hidden_size = 1024, + num_hidden_layers = 4, + num_attention_heads = 16, + num_key_value_heads = 4, + intermediate_size = 2048, + vocab_size = 32000, + tie_word_embeddings = False, + head_dim = 64, + layer_types = ["full_attention"] * 4, + vocab_size_per_layer_input = 256, + hidden_size_per_layer_input = 96, + ) + return replace(base, **overrides) + + def test_per_layer_input_increases_total_params(self): + with_ple = self._arch() + without_ple = replace(with_ple, hidden_size_per_layer_input = 0) + self.assertGreater( + compute_total_params(with_ple), + compute_total_params(without_ple), + ) + + def test_per_layer_input_modules_count_quantizable_block(self): + with_ple = self._arch() + without_ple = replace(with_ple, hidden_size_per_layer_input = 0) + # The PLE block adds: model_projection (hd*nl*pli), per_layer_input_gate + # (hd*pli per layer) + per_layer_projection (pli*hd per layer) as + # quantizable text linears. + n_layers = with_ple.num_hidden_layers + hd = with_ple.hidden_size + pli = with_ple.hidden_size_per_layer_input + expected_quantizable_extra = ( + hd * (n_layers * pli) + (hd * pli) * n_layers + (pli * hd) * n_layers + ) + delta = compute_total_params(with_ple) - compute_total_params(without_ple) + self.assertGreaterEqual(delta, expected_quantizable_extra) + + def test_all_linear_lora_excludes_per_layer_input_modules(self): + # why: Unsloth's get_peft_regex requires module names to contain a + # component tag (mlp/attn/...); PLE module names (per_layer_input_gate, + # per_layer_projection, per_layer_model_projection) lack any tag, so + # all-linear training does NOT attach LoRA to them. + arch = self._arch() + without_ple = replace(arch, hidden_size_per_layer_input = 0) + self.assertEqual( + compute_lora_params(arch, 16, ["all-linear"]), + compute_lora_params(without_ple, 16, ["all-linear"]), + ) + + def test_explicit_target_modules_does_not_add_per_layer_input(self): + arch = self._arch() + without_ple = replace(arch, hidden_size_per_layer_input = 0) + self.assertEqual( + compute_lora_params(arch, 16, ["q_proj", "v_proj"]), + compute_lora_params(without_ple, 16, ["q_proj", "v_proj"]), + ) + + +class TestDenseMlpLayerFallback(unittest.TestCase): + def test_falls_back_to_count_when_indices_empty(self): + from utils.hardware.vram_estimation import _is_dense_mlp_layer + + arch = ModelArchConfig( + hidden_size = 512, + num_hidden_layers = 4, + num_attention_heads = 8, + num_key_value_heads = 2, + intermediate_size = 1024, + vocab_size = 1024, + num_experts = 4, + moe_intermediate_size = 256, + num_dense_layers = 2, + ) + self.assertTrue(_is_dense_mlp_layer(arch, 0)) + self.assertTrue(_is_dense_mlp_layer(arch, 1)) + self.assertFalse(_is_dense_mlp_layer(arch, 2)) + self.assertFalse(_is_dense_mlp_layer(arch, 3)) + + +class TestExpertsSkipGranularity(unittest.TestCase): + def _arch(self): + return ModelArchConfig( + hidden_size = 512, + num_hidden_layers = 4, + num_attention_heads = 8, + num_key_value_heads = 2, + intermediate_size = 1024, + vocab_size = 1024, + tie_word_embeddings = True, + num_experts = 8, + moe_intermediate_size = 512, + num_dense_layers = 0, + head_dim = 64, + layer_types = ["full_attention"] * 4, + moe_has_dense_mlp = True, + ) + + def test_experts_skip_excludes_parallel_dense_projections(self): + no_skip = self._arch() + skip_experts = replace( + no_skip, + quantization_skip_modules = ["model.layers.0.mlp.experts"], + ) + skip_full_mlp = replace( + no_skip, + quantization_skip_modules = ["model.layers.0.mlp"], + ) + bytes_no_skip = compute_model_weights_bytes(no_skip, "qlora", True) + bytes_skip_experts = compute_model_weights_bytes(skip_experts, "qlora", True) + bytes_skip_mlp = compute_model_weights_bytes(skip_full_mlp, "qlora", True) + # why: under gemma4 enable_moe_block, `self.experts` is a sibling of + # `self.mlp`; skipping `model.layers.0.mlp` should cover only the + # dense MLP, while `model.layers.0.mlp.experts` covers the routed + # experts. Routed experts have far more params than the dense MLP, + # so skipping experts must add more bytes than skipping the dense + # path. + self.assertGreater(bytes_skip_experts, bytes_no_skip) + self.assertGreater(bytes_skip_mlp, bytes_no_skip) + self.assertGreater(bytes_skip_experts, bytes_skip_mlp) class TestSharedExperts(unittest.TestCase): @@ -608,6 +1325,16 @@ class TestMLA(unittest.TestCase): lora_p = compute_lora_params(DEEPSEEK_V3, 16, ["q_proj", "v_proj", "o_proj"]) self.assertGreater(lora_p, 0) + def test_mla_with_head_dim_does_not_route_through_structured(self): + from utils.hardware.vram_estimation import _uses_structured_layer_shapes + + mla_with_head_dim = replace(DEEPSEEK_V3, head_dim = 128) + self.assertFalse(_uses_structured_layer_shapes(mla_with_head_dim)) + self.assertEqual( + compute_lora_params(DEEPSEEK_V3, 16, ["q_proj", "v_proj", "o_proj"]), + compute_lora_params(mla_with_head_dim, 16, ["q_proj", "v_proj", "o_proj"]), + ) + class TestDenseMoEMix(unittest.TestCase): def test_dense_layers_change_total(self): @@ -691,5 +1418,952 @@ class TestDenseMoEMix(unittest.TestCase): self.assertNotEqual(lora_all, lora_mix) +class TestMlpLayerTypesDispatch(unittest.TestCase): + def _hf(self, **fields): + text_config = SimpleNamespace( + hidden_size = 64, + num_hidden_layers = 4, + num_attention_heads = 4, + num_key_value_heads = 4, + intermediate_size = 128, + vocab_size = 1000, + tie_word_embeddings = True, + num_local_experts = 4, + moe_intermediate_size = 32, + **fields, + ) + return SimpleNamespace(text_config = text_config, quantization_config = {}) + + def test_mlp_layer_types_drives_dense_indices(self): + hf = self._hf(mlp_layer_types = ["sparse", "dense", "sparse", "dense"]) + arch = extract_arch_config(hf) + self.assertIsNotNone(arch) + self.assertEqual(arch.dense_layer_indices, (1, 3)) + self.assertEqual(arch.num_dense_layers, 2) + + def test_mlp_layer_types_takes_priority_over_first_k_dense_replace(self): + hf = self._hf( + mlp_layer_types = ["dense", "sparse", "dense", "sparse"], + first_k_dense_replace = 3, + ) + arch = extract_arch_config(hf) + self.assertEqual(arch.dense_layer_indices, (0, 2)) + + def test_mlp_layer_types_ignores_unknown_entries(self): + hf = self._hf(mlp_layer_types = ["dense", "moe", "dense", "linear"]) + arch = extract_arch_config(hf) + self.assertEqual(arch.dense_layer_indices, (0, 2)) + + def test_mlp_layer_types_shorter_than_layers_only_marks_present(self): + hf = self._hf(mlp_layer_types = ["dense", "sparse"]) + arch = extract_arch_config(hf) + self.assertEqual(arch.dense_layer_indices, (0,)) + + def test_empty_mlp_layer_types_falls_through_to_first_k(self): + hf = self._hf(mlp_layer_types = [], first_k_dense_replace = 2) + arch = extract_arch_config(hf) + self.assertEqual(arch.dense_layer_indices, (0, 1)) + + +class TestPerLayerInputSkipAlias(unittest.TestCase): + def _hf(self, skip): + text_config = SimpleNamespace( + hidden_size = 64, + num_hidden_layers = 2, + num_attention_heads = 4, + num_key_value_heads = 4, + intermediate_size = 128, + vocab_size = 1000, + tie_word_embeddings = True, + hidden_size_per_layer_input = 8, + vocab_size_per_layer_input = 256, + ) + return SimpleNamespace( + text_config = text_config, + quantization_config = {"llm_int8_skip_modules": list(skip)}, + ) + + def test_per_layer_input_gate_skip_pulls_nonzero_delta(self): + from utils.hardware.vram_estimation import _compute_skipped_quantizable_elements + + arch = extract_arch_config(self._hf(["model.layers.0.per_layer_input_gate"])) + delta = _compute_skipped_quantizable_elements(arch) + self.assertEqual(delta, arch.hidden_size * arch.hidden_size_per_layer_input) + + def test_per_layer_model_projection_skip_pulls_global_delta(self): + from utils.hardware.vram_estimation import _compute_skipped_quantizable_elements + + arch = extract_arch_config(self._hf(["model.per_layer_model_projection"])) + delta = _compute_skipped_quantizable_elements(arch) + self.assertEqual( + delta, + arch.hidden_size + * arch.num_hidden_layers + * arch.hidden_size_per_layer_input, + ) + + def test_layer_aggregate_skip_includes_per_layer_input_modules(self): + from utils.hardware.vram_estimation import ( + _compute_skipped_quantizable_elements, + ) + + arch_with = extract_arch_config(self._hf(["model.layers.0"])) + # The text.layers.0 aggregate must include the PLE per-layer modules, + # so the same skip on a config without PLE produces a smaller value. + arch_without = extract_arch_config( + SimpleNamespace( + text_config = SimpleNamespace( + hidden_size = 64, + num_hidden_layers = 2, + num_attention_heads = 4, + num_key_value_heads = 4, + intermediate_size = 128, + vocab_size = 1000, + tie_word_embeddings = True, + hidden_size_per_layer_input = 0, + vocab_size_per_layer_input = 0, + ), + quantization_config = {"llm_int8_skip_modules": ["model.layers.0"]}, + ) + ) + self.assertGreater( + _compute_skipped_quantizable_elements(arch_with), + _compute_skipped_quantizable_elements(arch_without), + ) + + +class TestAllLinearStringHandling(unittest.TestCase): + def test_compute_lora_params_accepts_bare_all_linear_string(self): + list_form = compute_lora_params(LLAMA_8B, 16, ["all-linear"]) + str_form = compute_lora_params(LLAMA_8B, 16, "all-linear") + self.assertEqual(list_form, str_form) + self.assertGreater(list_form, 0) + + def test_compute_lora_params_string_with_underscores_normalized(self): + list_form = compute_lora_params(LLAMA_8B, 16, ["all_linear"]) + str_form = compute_lora_params(LLAMA_8B, 16, "all_linear") + self.assertEqual(list_form, str_form) + self.assertGreater(str_form, 0) + + +class TestSharedExpertVariants(unittest.TestCase): + def _hf(self, **fields): + text_config = SimpleNamespace( + hidden_size = 256, + num_hidden_layers = 4, + num_attention_heads = 8, + num_key_value_heads = 4, + intermediate_size = 1024, + vocab_size = 1000, + tie_word_embeddings = False, + num_local_experts = 8, + moe_intermediate_size = 128, + **fields, + ) + return SimpleNamespace(text_config = text_config, quantization_config = {}) + + def test_shared_expert_intermediate_size_extracted_and_infers_count(self): + arch = extract_arch_config(self._hf(shared_expert_intermediate_size = 64)) + self.assertEqual(arch.shared_expert_intermediate_size, 64) + self.assertEqual(arch.n_shared_experts, 1) + + def test_num_shared_experts_alias_extracted(self): + arch = extract_arch_config(self._hf(num_shared_experts = 2)) + self.assertEqual(arch.n_shared_experts, 2) + + def test_n_shared_experts_takes_priority_over_alias(self): + arch = extract_arch_config(self._hf(n_shared_experts = 3, num_shared_experts = 99)) + self.assertEqual(arch.n_shared_experts, 3) + + def test_shared_expert_size_separate_from_routed_changes_weight_count(self): + from utils.hardware.vram_estimation import _compute_moe_mlp_elements + + arch_separate = extract_arch_config( + self._hf(shared_expert_intermediate_size = 64) + ) + arch_implicit = extract_arch_config(self._hf(n_shared_experts = 1)) + # Different shared sizes (64 vs default moe_intermediate_size=128) must + # produce different MoE element counts. + self.assertNotEqual( + _compute_moe_mlp_elements(arch_separate), + _compute_moe_mlp_elements(arch_implicit), + ) + + def test_shared_expert_gate_counted_only_for_qwen_style(self): + from utils.hardware.vram_estimation import _compute_moe_mlp_elements + + # Qwen-style: shared_expert_intermediate_size set -> shared_expert_gate counted. + qwen_arch = extract_arch_config(self._hf(shared_expert_intermediate_size = 64)) + hd = qwen_arch.hidden_size + ms = qwen_arch.moe_intermediate_size + ne = qwen_arch.num_experts + ss = qwen_arch.shared_expert_intermediate_size + expected = hd * ms * 3 * ne + ne * hd + hd * ss * 3 * 1 + 1 * hd + self.assertEqual(_compute_moe_mlp_elements(qwen_arch), expected) + + # Non-Qwen shared experts (e.g. Exaone-MoE) -> no shared_expert_gate. + plain_arch = extract_arch_config(self._hf(n_shared_experts = 1)) + hd = plain_arch.hidden_size + ms = plain_arch.moe_intermediate_size + ne = plain_arch.num_experts + expected_plain = hd * ms * 3 * ne + ne * hd + hd * ms * 3 * 1 + self.assertEqual(_compute_moe_mlp_elements(plain_arch), expected_plain) + + +class TestSharedExpertActivation(unittest.TestCase): + def _make(self, **fields): + text_config = SimpleNamespace( + hidden_size = 512, + num_hidden_layers = 4, + num_attention_heads = 8, + num_key_value_heads = 4, + intermediate_size = 1024, + vocab_size = 1000, + tie_word_embeddings = False, + num_local_experts = 4, + moe_intermediate_size = 64, + **fields, + ) + return extract_arch_config( + SimpleNamespace(text_config = text_config, quantization_config = {}) + ) + + def test_shared_expert_increases_activation_bytes(self): + with_shared = self._make(shared_expert_intermediate_size = 64) + without = self._make() + self.assertGreater( + compute_activation_bytes( + with_shared, + 2, + 1024, + "none", + is_lora = True, + attention_implementation = "flash_attention_2", + ), + compute_activation_bytes( + without, + 2, + 1024, + "none", + is_lora = True, + attention_implementation = "flash_attention_2", + ), + ) + + def test_shared_expert_plus_dense_block_compose(self): + # gemma4 enable_moe_block with hypothetical shared expert: dense + routed + # + shared all live per layer; mlp_size should sum all three terms. + from utils.hardware.vram_estimation import _layer_qkv_mlp_sizes + + arch = self._make( + enable_moe_block = True, + shared_expert_intermediate_size = 32, + head_dim = 64, + layer_types = ["full_attention"] * 4, + ) + _, mlp_size = _layer_qkv_mlp_sizes(arch, 0) + # routed (64) + shared (32) + parallel dense intermediate (1024) + self.assertEqual(mlp_size, 64 + 32 + 1024) + + +class TestPerLayerInputActivation(unittest.TestCase): + def _make(self, **fields): + text_config = SimpleNamespace( + hidden_size = 512, + num_hidden_layers = 4, + num_attention_heads = 8, + num_key_value_heads = 4, + intermediate_size = 1024, + vocab_size = 1000, + tie_word_embeddings = False, + **fields, + ) + return extract_arch_config( + SimpleNamespace(text_config = text_config, quantization_config = {}) + ) + + def test_ple_increases_activation_bytes(self): + with_ple = self._make( + hidden_size_per_layer_input = 64, + vocab_size_per_layer_input = 256, + ) + without = self._make() + self.assertGreater( + compute_activation_bytes( + with_ple, + 2, + 1024, + "none", + is_lora = True, + attention_implementation = "flash_attention_2", + ), + compute_activation_bytes( + without, + 2, + 1024, + "none", + is_lora = True, + attention_implementation = "flash_attention_2", + ), + ) + + def test_ple_zero_does_not_inflate_activations(self): + without = self._make(hidden_size_per_layer_input = 0) + baseline = self._make() + self.assertEqual( + compute_activation_bytes( + without, + 2, + 512, + "none", + is_lora = True, + attention_implementation = "flash_attention_2", + ), + compute_activation_bytes( + baseline, + 2, + 512, + "none", + is_lora = True, + attention_implementation = "flash_attention_2", + ), + ) + + +class TestKvSharedActivation(unittest.TestCase): + def _make(self, kv_shared): + text_config = SimpleNamespace( + hidden_size = 512, + num_hidden_layers = 4, + num_attention_heads = 8, + num_key_value_heads = 4, + intermediate_size = 1024, + vocab_size = 1000, + tie_word_embeddings = False, + head_dim = 64, + num_kv_shared_layers = kv_shared, + layer_types = ["full_attention"] * 4, + ) + return extract_arch_config( + SimpleNamespace(text_config = text_config, quantization_config = {}) + ) + + def test_kv_shared_layers_keep_activation_bytes(self): + shared = self._make(kv_shared = 2) + full = self._make(kv_shared = 0) + self.assertEqual( + compute_activation_bytes( + shared, + 2, + 1024, + "none", + is_lora = True, + attention_implementation = "flash_attention_2", + ), + compute_activation_bytes( + full, + 2, + 1024, + "none", + is_lora = True, + attention_implementation = "flash_attention_2", + ), + ) + + +class TestSparseMoeSkipAliases(unittest.TestCase): + def _hf(self, skip, **fields): + text_config = SimpleNamespace( + hidden_size = 128, + num_hidden_layers = 2, + num_attention_heads = 4, + num_key_value_heads = 4, + intermediate_size = 256, + vocab_size = 1000, + tie_word_embeddings = False, + num_local_experts = 4, + moe_intermediate_size = 64, + **fields, + ) + return SimpleNamespace( + text_config = text_config, + quantization_config = {"llm_int8_skip_modules": list(skip)}, + ) + + def test_gemma4_layers_experts_alias_pulls_routed(self): + from utils.hardware.vram_estimation import _compute_skipped_quantizable_elements + + arch = extract_arch_config( + self._hf(["model.layers.0.experts"], enable_moe_block = True) + ) + self.assertGreater(_compute_skipped_quantizable_elements(arch), 0) + + def test_qwen_shared_expert_skip_pulls_only_shared(self): + from utils.hardware.vram_estimation import _compute_skipped_quantizable_elements + + arch = extract_arch_config( + self._hf( + ["model.layers.0.mlp.shared_expert"], + shared_expert_intermediate_size = 32, + ) + ) + # shared_expert delta only -- routed mlp.experts is NOT skipped. + delta = _compute_skipped_quantizable_elements(arch) + self.assertGreater(delta, 0) + full_layer = extract_arch_config( + self._hf( + ["model.layers.0.mlp"], + shared_expert_intermediate_size = 32, + ) + ) + self.assertGreater( + _compute_skipped_quantizable_elements(full_layer), + delta, + ) + + def test_exaone_shared_experts_plural_alias(self): + from utils.hardware.vram_estimation import _compute_skipped_quantizable_elements + + arch = extract_arch_config( + self._hf( + ["model.layers.0.mlp.shared_experts"], + num_shared_experts = 1, + ) + ) + self.assertGreater(_compute_skipped_quantizable_elements(arch), 0) + + +class TestAllLinearMoELoraExclusion(unittest.TestCase): + def _arch(self, **fields): + text_config = SimpleNamespace( + hidden_size = 256, + num_hidden_layers = 2, + num_attention_heads = 4, + num_key_value_heads = 4, + intermediate_size = 512, + vocab_size = 1000, + tie_word_embeddings = False, + num_local_experts = 8, + moe_intermediate_size = 64, + **fields, + ) + return extract_arch_config( + SimpleNamespace(text_config = text_config, quantization_config = {}) + ) + + def test_all_linear_drops_routed_moe_expert_lora(self): + arch = self._arch() + all_linear = compute_lora_params(arch, 8, "all-linear") + explicit = compute_lora_params(arch, 8, ["gate_proj", "up_proj", "down_proj"]) + self.assertLess(all_linear, explicit) + + def test_all_linear_drops_shared_expert_lora(self): + arch = self._arch(shared_expert_intermediate_size = 32) + all_linear = compute_lora_params(arch, 8, "all-linear") + explicit = compute_lora_params(arch, 8, ["gate_proj", "up_proj", "down_proj"]) + # explicit includes routed + shared MoE; all-linear includes neither. + self.assertLess(all_linear, explicit) + + def test_all_linear_includes_attention_lora(self): + arch = self._arch() + all_linear = compute_lora_params(arch, 8, "all-linear") + attn_only = compute_lora_params( + arch, 8, ["q_proj", "k_proj", "v_proj", "o_proj"] + ) + # all-linear still attaches to attention nn.Linear modules. + self.assertGreaterEqual(all_linear, attn_only) + + +class TestExplicitPerLayerInputLora(unittest.TestCase): + def _arch(self): + text_config = SimpleNamespace( + hidden_size = 256, + num_hidden_layers = 3, + num_attention_heads = 4, + num_key_value_heads = 4, + intermediate_size = 512, + vocab_size = 1000, + tie_word_embeddings = False, + hidden_size_per_layer_input = 32, + vocab_size_per_layer_input = 128, + ) + return extract_arch_config( + SimpleNamespace(text_config = text_config, quantization_config = {}) + ) + + def test_explicit_per_layer_input_gate_returns_nonzero(self): + arch = self._arch() + result = compute_lora_params(arch, 16, ["per_layer_input_gate"]) + self.assertGreater(result, 0) + + def test_explicit_per_layer_projection_returns_nonzero(self): + arch = self._arch() + result = compute_lora_params(arch, 16, ["per_layer_projection"]) + self.assertGreater(result, 0) + + def test_explicit_per_layer_model_projection_returns_nonzero(self): + arch = self._arch() + result = compute_lora_params(arch, 16, ["per_layer_model_projection"]) + self.assertGreater(result, 0) + + def test_explicit_ple_string_target_handled(self): + # Bare-string target with a PLE name should not be iterated char-by-char. + arch = self._arch() + list_form = compute_lora_params(arch, 16, ["per_layer_input_gate"]) + str_form = compute_lora_params(arch, 16, "per_layer_input_gate") + self.assertEqual(list_form, str_form) + + +class TestTopKExpertActivation(unittest.TestCase): + def _make(self, **fields): + text_config = SimpleNamespace( + hidden_size = 512, + num_hidden_layers = 4, + num_attention_heads = 8, + num_key_value_heads = 4, + intermediate_size = 1024, + vocab_size = 1000, + tie_word_embeddings = False, + num_local_experts = 8, + moe_intermediate_size = 64, + **fields, + ) + return extract_arch_config( + SimpleNamespace(text_config = text_config, quantization_config = {}) + ) + + def test_num_experts_per_tok_extracted(self): + arch = self._make(num_experts_per_tok = 4) + self.assertEqual(arch.num_experts_per_tok, 4) + + def test_top_k_experts_alias_extracted(self): + arch = self._make(top_k_experts = 8) + self.assertEqual(arch.num_experts_per_tok, 8) + + def test_default_top_k_one_unchanged(self): + arch = self._make() + self.assertEqual(arch.num_experts_per_tok, 1) + + def test_top_k_scales_moe_activation(self): + single = self._make() + multi = self._make(num_experts_per_tok = 8) + single_act = compute_activation_bytes( + single, + 2, + 512, + "none", + is_lora = True, + attention_implementation = "flash_attention_2", + ) + multi_act = compute_activation_bytes( + multi, + 2, + 512, + "none", + is_lora = True, + attention_implementation = "flash_attention_2", + ) + self.assertGreater(multi_act, single_act) + + +class TestErnieMoEListConfig(unittest.TestCase): + def _hf(self, **fields): + text_config = SimpleNamespace( + hidden_size = 256, + num_hidden_layers = 4, + num_attention_heads = 4, + num_key_value_heads = 4, + intermediate_size = 1024, + vocab_size = 1000, + tie_word_embeddings = False, + **fields, + ) + return SimpleNamespace(text_config = text_config, quantization_config = {}) + + def test_list_moe_intermediate_size_scalarized(self): + arch = extract_arch_config( + self._hf( + moe_num_experts = 32, + moe_intermediate_size = [1536, 512], + ) + ) + # why: ERNIE 4.5 VL MoE encodes [text_routed, vision_routed]; the + # second element is the vision-routed expert width, not the shared + # expert width. Shared experts are sized from the text-routed width + # (= moe_intermediate_size[0]) when moe_num_shared_experts is set. + self.assertEqual(arch.moe_intermediate_size, 1536) + self.assertIsNone(arch.shared_expert_intermediate_size) + self.assertEqual(arch.n_shared_experts, 0) + + def test_moe_num_experts_alias_extracted(self): + arch = extract_arch_config( + self._hf( + moe_num_experts = 64, + moe_intermediate_size = 1024, + ) + ) + self.assertEqual(arch.num_experts, 64) + + def test_moe_num_shared_experts_alias_extracted(self): + arch = extract_arch_config( + self._hf( + moe_num_experts = 16, + moe_num_shared_experts = 2, + moe_intermediate_size = 1024, + ) + ) + self.assertEqual(arch.n_shared_experts, 2) + + def test_explicit_shared_size_overrides_list_second_element(self): + arch = extract_arch_config( + self._hf( + moe_num_experts = 8, + moe_intermediate_size = [1536, 512], + shared_expert_intermediate_size = 256, + ) + ) + # Explicit shared size wins over moe_intermediate_size[1]. + self.assertEqual(arch.shared_expert_intermediate_size, 256) + + +class TestSuffixSkipModuleMatch(unittest.TestCase): + def _hf(self, skip): + text_config = SimpleNamespace( + hidden_size = 128, + num_hidden_layers = 2, + num_attention_heads = 4, + num_key_value_heads = 4, + intermediate_size = 256, + vocab_size = 1000, + tie_word_embeddings = False, + ) + return SimpleNamespace( + text_config = text_config, + quantization_config = {"llm_int8_skip_modules": list(skip)}, + ) + + def test_q_proj_suffix_skip_matches_all_layers(self): + from utils.hardware.vram_estimation import _compute_skipped_quantizable_elements + + arch = extract_arch_config(self._hf(["q_proj"])) + delta = _compute_skipped_quantizable_elements(arch) + # 2 layers * hd * hd of q_proj weight elements. + self.assertEqual(delta, 2 * arch.hidden_size * arch.hidden_size) + + def test_self_attn_aggregate_skip_matches_aggregate(self): + from utils.hardware.vram_estimation import _compute_skipped_quantizable_elements + + arch = extract_arch_config(self._hf(["self_attn"])) + # The aggregate text.layers..self_attn matches; total covers both layers. + delta = _compute_skipped_quantizable_elements(arch) + self.assertGreater(delta, 0) + + def test_vision_prefix_skip_does_not_match_text_alias(self): + from utils.hardware.vram_estimation import _module_path_matches + + # vision_tower-prefixed full path must NOT match text-tower aliases. + self.assertFalse( + _module_path_matches( + "vision_tower.model.layers.0.self_attn.q_proj", + "model.layers.0.self_attn.q_proj", + ) + ) + + +class TestMultimodalFullModelBytes(unittest.TestCase): + def test_extra_bytes_added_when_safetensors_exceeds_text_arch(self): + from utils.hardware import hardware as hardware_module + + config = SimpleNamespace( + hidden_size = 1024, + num_hidden_layers = 4, + num_attention_heads = 8, + num_key_value_heads = 4, + intermediate_size = 2048, + vocab_size = 32000, + tie_word_embeddings = False, + ) + # Force safetensors size >>> arch text-only bytes. + big_safetensors = 20 * 1024**3 + with ( + patch.object( + hardware_module, + "_load_config_for_gpu_estimate", + return_value = config, + ), + patch.object( + hardware_module, + "estimate_fp16_model_size_bytes", + return_value = (big_safetensors, "safetensors"), + ), + patch.object( + hardware_module, + "_determine_attention_impl_for_gpu_estimate", + return_value = "flash_attention_2", + ), + patch.object( + hardware_module, + "get_visible_gpu_count", + return_value = 1, + ), + ): + _, metadata = hardware_module.estimate_required_model_memory_gb( + "fake/model", + training_type = "LoRA/QLoRA", + load_in_4bit = True, + ) + self.assertEqual(metadata.get("estimation_mode"), "detailed") + # model_weights_gb must reflect the extra non-text bytes (>5 GB + # since text-only arch_fp16 is small for these dims). + self.assertGreater(metadata["vram_breakdown"]["model_weights_gb"], 5.0) + + def test_no_extra_when_safetensors_smaller_than_text_arch(self): + from utils.hardware import hardware as hardware_module + + config = SimpleNamespace( + hidden_size = 4096, + num_hidden_layers = 32, + num_attention_heads = 32, + num_key_value_heads = 8, + intermediate_size = 11008, + vocab_size = 32000, + tie_word_embeddings = False, + ) + tiny_safetensors = 100 # bytes, deliberately absurdly small + with ( + patch.object( + hardware_module, + "_load_config_for_gpu_estimate", + return_value = config, + ), + patch.object( + hardware_module, + "estimate_fp16_model_size_bytes", + return_value = (tiny_safetensors, "safetensors"), + ), + patch.object( + hardware_module, + "_determine_attention_impl_for_gpu_estimate", + return_value = "flash_attention_2", + ), + patch.object( + hardware_module, + "get_visible_gpu_count", + return_value = 1, + ), + ): + required, metadata = hardware_module.estimate_required_model_memory_gb( + "fake/model", + training_type = "LoRA/QLoRA", + load_in_4bit = True, + ) + # No negative extra; required_gb stays a positive finite number. + self.assertGreater(required, 0) + + +class TestLlama4ArchExtraction(unittest.TestCase): + def _llama4_text_config(self, **fields): + base = dict( + hidden_size = 2048, + num_hidden_layers = 4, + num_attention_heads = 16, + num_key_value_heads = 4, + intermediate_size = 8192, + intermediate_size_mlp = 16384, + vocab_size = 32000, + tie_word_embeddings = True, + num_local_experts = 4, + num_experts_per_tok = 2, + ) + base.update(fields) + return SimpleNamespace(**base) + + def test_llama4_moe_layers_dispatch_uses_explicit_indices(self): + from utils.hardware.vram_estimation import _compute_dense_layer_indices + + cfg = SimpleNamespace(num_hidden_layers = 4, moe_layers = [1, 3]) + self.assertEqual(_compute_dense_layer_indices(cfg, 4), (0, 2)) + + def test_llama4_moe_layers_takes_priority_over_first_k_dense_replace(self): + from utils.hardware.vram_estimation import _compute_dense_layer_indices + + cfg = SimpleNamespace( + num_hidden_layers = 6, + moe_layers = [2, 4], + first_k_dense_replace = 4, + ) + self.assertEqual(_compute_dense_layer_indices(cfg, 6), (0, 1, 3, 5)) + + def test_dense_intermediate_size_picks_up_intermediate_size_mlp(self): + from utils.hardware.vram_estimation import _dense_mlp_size + + arch = extract_arch_config(self._llama4_text_config(moe_layers = [1, 3])) + self.assertIsNotNone(arch) + self.assertEqual(arch.intermediate_size, 8192) + self.assertEqual(arch.dense_intermediate_size, 16384) + self.assertEqual(_dense_mlp_size(arch), 16384) + + def test_auto_attaches_one_shared_expert_at_routed_width(self): + from utils.hardware.vram_estimation import _shared_expert_size + + arch = extract_arch_config(self._llama4_text_config(moe_layers = [1, 3])) + self.assertIsNotNone(arch) + self.assertEqual(arch.n_shared_experts, 1) + self.assertIsNone(arch.shared_expert_intermediate_size) + self.assertEqual(_shared_expert_size(arch), arch.intermediate_size) + + def test_non_llama4_config_leaves_dense_intermediate_size_none(self): + from utils.hardware.vram_estimation import _dense_mlp_size + + cfg = SimpleNamespace( + hidden_size = 1024, + num_hidden_layers = 4, + num_attention_heads = 8, + num_key_value_heads = 2, + intermediate_size = 4096, + vocab_size = 32000, + tie_word_embeddings = True, + ) + arch = extract_arch_config(cfg) + self.assertIsNotNone(arch) + self.assertIsNone(arch.dense_intermediate_size) + self.assertEqual(_dense_mlp_size(arch), 4096) + + def test_intermediate_size_mlp_without_moe_does_not_force_shared_expert(self): + cfg = SimpleNamespace( + hidden_size = 2048, + num_hidden_layers = 4, + num_attention_heads = 16, + num_key_value_heads = 4, + intermediate_size = 8192, + intermediate_size_mlp = 16384, + vocab_size = 32000, + tie_word_embeddings = True, + ) + arch = extract_arch_config(cfg) + self.assertIsNotNone(arch) + self.assertEqual(arch.dense_intermediate_size, 16384) + self.assertEqual(arch.n_shared_experts, 0) + + +class TestDbrxFfnConfigExtraction(unittest.TestCase): + def test_extracts_moe_fields_from_ffn_subconfig(self): + ffn = SimpleNamespace(moe_num_experts = 4, moe_top_k = 2, ffn_hidden_size = 1024) + cfg = SimpleNamespace( + hidden_size = 2048, + num_hidden_layers = 4, + num_attention_heads = 16, + num_key_value_heads = 4, + intermediate_size = 2048, + vocab_size = 32000, + tie_word_embeddings = False, + ffn_config = ffn, + ) + arch = extract_arch_config(cfg) + self.assertIsNotNone(arch) + self.assertEqual(arch.num_experts, 4) + self.assertEqual(arch.num_experts_per_tok, 2) + self.assertEqual(arch.moe_intermediate_size, 1024) + + def test_top_level_attrs_take_precedence_over_ffn_config(self): + ffn = SimpleNamespace(moe_num_experts = 4, moe_top_k = 2, ffn_hidden_size = 1024) + cfg = SimpleNamespace( + hidden_size = 2048, + num_hidden_layers = 4, + num_attention_heads = 16, + num_key_value_heads = 4, + intermediate_size = 2048, + vocab_size = 32000, + tie_word_embeddings = False, + ffn_config = ffn, + num_local_experts = 16, + num_experts_per_tok = 8, + ) + arch = extract_arch_config(cfg) + self.assertIsNotNone(arch) + self.assertEqual(arch.num_experts, 16) + self.assertEqual(arch.num_experts_per_tok, 8) + + +class TestErniePhaseModuloDispatch(unittest.TestCase): + def test_phase_modulo_with_interval_two_matches_decoder(self): + from utils.hardware.vram_estimation import _compute_dense_layer_indices + + cfg = SimpleNamespace( + num_hidden_layers = 10, + moe_layer_start_index = 2, + moe_layer_end_index = 8, + moe_layer_interval = 2, + ) + # Decoder gates by ((i + 1) % 2 == 0) AND 2 <= i <= 8 -> MoE = {3, 5, 7}. + self.assertEqual(_compute_dense_layer_indices(cfg, 10), (0, 1, 2, 4, 6, 8, 9)) + + def test_phase_modulo_with_interval_three(self): + from utils.hardware.vram_estimation import _compute_dense_layer_indices + + cfg = SimpleNamespace( + num_hidden_layers = 9, + moe_layer_start_index = 0, + moe_layer_end_index = -1, + moe_layer_interval = 3, + ) + self.assertEqual(_compute_dense_layer_indices(cfg, 9), (0, 1, 3, 4, 6, 7)) + + +class TestErnieVlSharedExpertWidth(unittest.TestCase): + def test_shared_expert_width_uses_text_routed_not_vision(self): + from utils.hardware.vram_estimation import ( + _compute_shared_moe_elements, + _shared_expert_size, + ) + + cfg = SimpleNamespace( + text_config = SimpleNamespace( + hidden_size = 1024, + num_hidden_layers = 4, + num_attention_heads = 8, + num_key_value_heads = 4, + intermediate_size = 2048, + vocab_size = 32000, + tie_word_embeddings = False, + moe_num_experts = 8, + moe_num_shared_experts = 2, + moe_intermediate_size = [1536, 512], + ), + quantization_config = {}, + ) + arch = extract_arch_config(cfg) + self.assertIsNotNone(arch) + self.assertIsNone(arch.shared_expert_intermediate_size) + self.assertEqual(arch.moe_intermediate_size, 1536) + self.assertEqual(arch.n_shared_experts, 2) + self.assertEqual(_shared_expert_size(arch), 1536) + self.assertEqual(_compute_shared_moe_elements(arch), 1024 * 1536 * 3 * 2) + + def test_qwen_style_explicit_shared_expert_size_still_adds_gate(self): + from utils.hardware.vram_estimation import _compute_shared_moe_elements + + cfg = SimpleNamespace( + hidden_size = 1024, + num_hidden_layers = 4, + num_attention_heads = 8, + num_key_value_heads = 4, + intermediate_size = 2048, + vocab_size = 32000, + tie_word_embeddings = False, + num_local_experts = 8, + moe_intermediate_size = 256, + shared_expert_intermediate_size = 768, + ) + arch = extract_arch_config(cfg) + self.assertIsNotNone(arch) + self.assertEqual(arch.shared_expert_intermediate_size, 768) + self.assertEqual(arch.n_shared_experts, 1) + self.assertEqual( + _compute_shared_moe_elements(arch), + 1024 * 768 * 3 + 1 * 1024, + ) + + if __name__ == "__main__": unittest.main() diff --git a/studio/backend/utils/hardware/VRAM_ESTIMATION.md b/studio/backend/utils/hardware/VRAM_ESTIMATION.md index 26072b208f..a6b4de29d2 100644 --- a/studio/backend/utils/hardware/VRAM_ESTIMATION.md +++ b/studio/backend/utils/hardware/VRAM_ESTIMATION.md @@ -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 ``` diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index be31c00a78..c218b7b4b9 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -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): diff --git a/studio/backend/utils/hardware/vram_estimation.py b/studio/backend/utils/hardware/vram_estimation.py index e03665374d..ba1b1dfe61 100644 --- a/studio/backend/utils/hardware/vram_estimation.py +++ b/studio/backend/utils/hardware/vram_estimation.py @@ -16,7 +16,26 @@ from dataclasses import dataclass, field from typing import Dict, Optional QUANT_4BIT_FACTOR = 16 / 5 +DOUBLE_QUANT_4BIT_FACTOR = ( + 3.6 # bnb_4bit_use_double_quant; see VRAM_ESTIMATION.md section 1 +) CUDA_OVERHEAD_BYTES = int(1.4 * 1024**3) # calibrated on RTX 5070 Ti +NON_FLASH_ATTENTION_FACTOR = ( + 12.0 # eager attention score+workspace overhead; see VRAM_ESTIMATION.md section 5 +) + +LINEAR_ATTENTION_IMPLS = frozenset({"flash_attention_2", "sdpa", "flex_attention"}) + +_SKIP_MODULE_TEXT_PREFIXES = frozenset( + { + "model", + "model.model", + "language_model", + "language_model.model", + "model.language_model", + "model.language_model.model", + } +) DEFAULT_TARGET_MODULES = [ "q_proj", @@ -27,6 +46,8 @@ DEFAULT_TARGET_MODULES = [ "up_proj", "down_proj", ] +ATTENTION_TARGET_MODULES = {"q_proj", "k_proj", "v_proj", "o_proj"} +MLP_TARGET_MODULES = {"gate_proj", "up_proj", "down_proj"} # Empirically calibrated bytes/param — see VRAM_ESTIMATION.md for rationale. OPTIMIZER_BYTES_PER_PARAM: Dict[str, int] = { @@ -61,12 +82,28 @@ class ModelArchConfig: num_experts: Optional[int] = None moe_intermediate_size: Optional[int] = None n_shared_experts: int = 0 + shared_expert_intermediate_size: Optional[int] = None + num_experts_per_tok: int = 1 num_dense_layers: int = 0 q_lora_rank: Optional[int] = None kv_lora_rank: Optional[int] = None qk_nope_head_dim: Optional[int] = None qk_rope_head_dim: Optional[int] = None v_head_dim: Optional[int] = None + head_dim: Optional[int] = None + global_head_dim: Optional[int] = None + num_global_key_value_heads: Optional[int] = None + attention_k_eq_v: bool = False + layer_types: Optional[list] = None + num_kv_shared_layers: int = 0 + use_double_wide_mlp: bool = False + vocab_size_per_layer_input: int = 0 + hidden_size_per_layer_input: int = 0 + quantization_skip_modules: list = field(default_factory = list) + quant_4bit_factor: float = QUANT_4BIT_FACTOR + moe_has_dense_mlp: bool = False + dense_layer_indices: tuple = () + dense_intermediate_size: Optional[int] = None @dataclass @@ -79,6 +116,7 @@ class TrainingVramConfig: gradient_checkpointing: str = "unsloth" optimizer: str = "adamw_8bit" load_in_4bit: bool = True + attention_implementation: str = "flash_attention_2" @dataclass @@ -89,8 +127,8 @@ class VramBreakdown: gradients: int activations: int cuda_overhead: int - # The computed (formula-based) activation cost before floors. - # This is the true per-layer cost that doesn't shard across GPUs. + # Equals `activations`; retained for backward compatibility with + # consumers that read this field. activations_computed: int = 0 @property @@ -108,17 +146,15 @@ class VramBreakdown: """Minimum VRAM a single GPU needs: its shard + non-shardable costs. Weights/LoRA/optimizer/gradients shard across GPUs. - The computed activation cost does NOT shard (one GPU runs the layer). - The floor portion (activations - computed) is overhead that shards. + Activations do NOT shard (the GPU running a layer holds them). """ shardable = ( self.model_weights + self.lora_adapters + self.optimizer_states + self.gradients - + (self.activations - self.activations_computed) # floor overhead shards ) - per_gpu_fixed = self.activations_computed + self.cuda_overhead + per_gpu_fixed = self.activations + self.cuda_overhead return shardable // max(n_gpus, 1) + per_gpu_fixed def to_gb_dict(self) -> Dict[str, float]: @@ -133,28 +169,88 @@ class VramBreakdown: } -def _compute_num_dense_layers(text_config, total_layers: int) -> int: - """Count how many layers use dense MLP instead of MoE.""" +def _first_scalar(value): + # why: ERNIE MoE configs ship moe_intermediate_size / moe_num_experts as + # [routed, shared] lists; downstream arithmetic needs the routed scalar. + if isinstance(value, (list, tuple)): + return value[0] if value else None + return value + + +def _max_scalar(value): + # why: Hunyuan-V1-MoE moe_topk can be a per-layer list; activation + # accounting uses the max top-k as a conservative upper bound. + if isinstance(value, (list, tuple)): + items = [v for v in value if v is not None] + return max(items) if items else None + return value + + +def _compute_dense_layer_indices(text_config, total_layers: int) -> tuple: + """Layer indices that use dense MLP instead of MoE. Position matters.""" + # why: transformers Exaone-MoE / Laguna / Hy_v3 / GLM-MoE-DSA / GLM4-MoE-Lite / + # Ernie4_5_VL_MoE prefer per-position `mlp_layer_types` over the prefix-style + # `first_k_dense_replace` and may omit `decoder_sparse_step` entirely. + layer_types = getattr(text_config, "mlp_layer_types", None) + if layer_types: + return tuple( + i + for i, t in enumerate(layer_types[:total_layers]) + if str(t).lower() == "dense" + ) + + # why: Llama4TextConfig.__init__ auto-populates self.moe_layers from + # interleave_moe_layer_step; Llama4TextDecoderLayer dispatches via + # `layer_idx in config.moe_layers` (modeling_llama4.py). + llama4_moe_layers = getattr(text_config, "moe_layers", None) + if llama4_moe_layers is not None: + moe_indices = {int(i) for i in llama4_moe_layers} + return tuple(i for i in range(total_layers) if i not in moe_indices) + + # why: transformers ERNIE 4.5 MoE / ERNIE 4.5 VL MoE declare MoE layers + # via moe_layer_start_index / moe_layer_end_index / moe_layer_interval; + # the model's per-layer guard is `(layer_idx + 1) % interval == 0` with + # start <= layer_idx <= end (modeling_ernie4_5_moe.py). + moe_start = getattr(text_config, "moe_layer_start_index", None) + moe_interval = getattr(text_config, "moe_layer_interval", None) + if moe_start is not None and moe_interval is not None and int(moe_interval) > 0: + moe_end_raw = getattr(text_config, "moe_layer_end_index", None) + end = ( + total_layers + if moe_end_raw is None or int(moe_end_raw) == -1 + else min(int(moe_end_raw) + 1, total_layers) + ) + start = max(0, int(moe_start)) + interval = int(moe_interval) + moe_indices = {i for i in range(start, end) if (i + 1) % interval == 0} + return tuple(i for i in range(total_layers) if i not in moe_indices) + first_k = getattr(text_config, "first_k_dense_replace", None) if first_k is not None: - return min(int(first_k), total_layers) + return tuple(range(min(int(first_k), total_layers))) sparse_step = getattr(text_config, "decoder_sparse_step", None) mlp_only = getattr(text_config, "mlp_only_layers", None) or [] if sparse_step is not None and sparse_step > 0: - mlp_only_set = set(mlp_only) - moe_count = sum( - 1 + mlp_only_set = {int(i) for i in mlp_only} + return tuple( + i for i in range(total_layers) - if i not in mlp_only_set and (i + 1) % sparse_step == 0 + if i in mlp_only_set or (i + 1) % sparse_step != 0 ) - return total_layers - moe_count - - return 0 + return () def extract_arch_config(hf_config) -> Optional[ModelArchConfig]: text_config = getattr(hf_config, "text_config", None) or hf_config + quantization_config = getattr(hf_config, "quantization_config", None) or {} + if not isinstance(quantization_config, dict): + quantization_config = getattr(quantization_config, "to_dict", lambda: {})() + quant_4bit_factor = ( + DOUBLE_QUANT_4BIT_FACTOR + if quantization_config.get("bnb_4bit_use_double_quant", False) + else QUANT_4BIT_FACTOR + ) hidden_size = getattr(text_config, "hidden_size", None) num_layers = getattr(text_config, "num_hidden_layers", None) @@ -177,18 +273,75 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]: num_kv_heads = getattr(text_config, "num_key_value_heads", num_heads) + # why: DBRX places its MoE attrs on the DbrxFFNConfig sub-config; probe + # ffn_config as a secondary source so DBRX is not misclassified as dense. + ffn_config = getattr(text_config, "ffn_config", None) + + def _moe_attr(name): + value = getattr(text_config, name, None) + if value is None and ffn_config is not None: + value = getattr(ffn_config, name, None) + return value + num_experts = None - for attr in ("num_local_experts", "num_experts", "n_routed_experts"): - num_experts = getattr(text_config, attr, None) + for attr in ( + "num_local_experts", + "num_experts", + "n_routed_experts", + "moe_num_experts", + ): + num_experts = _first_scalar(_moe_attr(attr)) if num_experts is not None: break - moe_intermediate = getattr(text_config, "moe_intermediate_size", None) - n_shared_experts = getattr(text_config, "n_shared_experts", None) or 0 + moe_intermediate_raw = _moe_attr("moe_intermediate_size") + if moe_intermediate_raw is None: + moe_intermediate_raw = _moe_attr("ffn_hidden_size") + moe_intermediate = _first_scalar(moe_intermediate_raw) + # why: Exaone-MoE / ERNIE families alias num_shared_experts / + # moe_num_shared_experts to the canonical n_shared_experts. + n_shared_experts = ( + _first_scalar(_moe_attr("n_shared_experts")) + or _first_scalar(_moe_attr("num_shared_experts")) + or _first_scalar(_moe_attr("moe_num_shared_experts")) + or 0 + ) + shared_expert_intermediate_size = _moe_attr("shared_expert_intermediate_size") + if shared_expert_intermediate_size and n_shared_experts == 0: + n_shared_experts = 1 + # why: DBRX exposes moe_top_k, Hunyuan-V1-MoE exposes moe_topk (which can + # be a per-layer list); _max_scalar normalizes list values to the worst + # case so int(...) below cannot crash on the canonical attribute_map path. + num_experts_per_tok = ( + _max_scalar(_moe_attr("num_experts_per_tok")) + or _max_scalar(_moe_attr("top_k_experts")) + or _max_scalar(_moe_attr("moe_top_k")) + or _max_scalar(_moe_attr("moe_topk")) + or 1 + ) - num_dense_layers = 0 + dense_layer_indices: tuple = () if num_experts is not None and num_experts > 1: - num_dense_layers = _compute_num_dense_layers(text_config, num_layers) + dense_layer_indices = _compute_dense_layer_indices(text_config, num_layers) + num_dense_layers = len(dense_layer_indices) + + # why: Llama4 dense layers use intermediate_size_mlp; routed and shared + # experts use intermediate_size. Llama4TextMoe builds one shared_expert + # per MoE layer (modeling_llama4.py). + intermediate_size_mlp_raw = _first_scalar(_moe_attr("intermediate_size_mlp")) + dense_intermediate_size = ( + int(intermediate_size_mlp_raw) + if intermediate_size_mlp_raw is not None + else None + ) + if ( + intermediate_size_mlp_raw is not None + and num_experts is not None + and num_experts > 1 + and shared_expert_intermediate_size is None + and n_shared_experts == 0 + ): + n_shared_experts = 1 q_lora_rank = getattr(text_config, "q_lora_rank", None) kv_lora_rank = getattr(text_config, "kv_lora_rank", None) @@ -207,15 +360,418 @@ def extract_arch_config(hf_config) -> Optional[ModelArchConfig]: num_experts = num_experts, moe_intermediate_size = moe_intermediate, n_shared_experts = n_shared_experts, + shared_expert_intermediate_size = shared_expert_intermediate_size, + num_experts_per_tok = int(num_experts_per_tok), num_dense_layers = num_dense_layers, q_lora_rank = q_lora_rank, kv_lora_rank = kv_lora_rank, qk_nope_head_dim = qk_nope_head_dim, qk_rope_head_dim = qk_rope_head_dim, v_head_dim = v_head_dim, + head_dim = getattr(text_config, "head_dim", None), + global_head_dim = getattr(text_config, "global_head_dim", None), + num_global_key_value_heads = getattr( + text_config, + "num_global_key_value_heads", + None, + ), + attention_k_eq_v = bool(getattr(text_config, "attention_k_eq_v", False)), + layer_types = getattr(text_config, "layer_types", None), + num_kv_shared_layers = getattr(text_config, "num_kv_shared_layers", None) or 0, + use_double_wide_mlp = bool(getattr(text_config, "use_double_wide_mlp", False)), + vocab_size_per_layer_input = getattr( + text_config, + "vocab_size_per_layer_input", + None, + ) + or 0, + hidden_size_per_layer_input = getattr( + text_config, + "hidden_size_per_layer_input", + None, + ) + or 0, + quantization_skip_modules = list( + quantization_config.get("llm_int8_skip_modules", []) or [] + ), + quant_4bit_factor = quant_4bit_factor, + moe_has_dense_mlp = bool(getattr(text_config, "enable_moe_block", False)), + dense_layer_indices = dense_layer_indices, + dense_intermediate_size = dense_intermediate_size, ) +def _targets_all_linear(target_modules) -> bool: + # why: peft LoraConfig accepts target_modules="all-linear" as a bare + # string; iterating a string yields chars and never matches the set. + if isinstance(target_modules, str): + target_modules = [target_modules] + normalized = {str(module).lower().replace("_", "-") for module in target_modules} + return normalized == {"all-linear"} + + +def _head_dim(arch: ModelArchConfig) -> int: + return arch.head_dim or arch.hidden_size // arch.num_attention_heads + + +def _layer_types(arch: ModelArchConfig) -> list: + if arch.layer_types and len(arch.layer_types) == arch.num_hidden_layers: + return arch.layer_types + return ["full_attention"] * arch.num_hidden_layers + + +def _uses_structured_layer_shapes(arch: ModelArchConfig) -> bool: + # MLA configs have their own q/kv low-rank projection shape formulas in + # _compute_attn_elements / _lora_attn_elements; do not let head_dim or + # other structured fields override that path. + if arch.q_lora_rank is not None: + return False + return bool( + arch.layer_types + or arch.head_dim is not None + or arch.global_head_dim is not None + or arch.num_global_key_value_heads is not None + or arch.attention_k_eq_v + or arch.num_kv_shared_layers > 0 + or arch.use_double_wide_mlp + ) + + +def _is_kv_shared_layer(arch: ModelArchConfig, layer_idx: int) -> bool: + if arch.num_kv_shared_layers <= 0: + return False + first_shared = arch.num_hidden_layers - arch.num_kv_shared_layers + # why: transformers Gemma4 (modeling_gemma4.py:1031, modular_gemma4.py:863) + # uses the same `> 0` guard so a fully-shared config raises during model + # construction; matching upstream avoids producing a detailed estimate + # for a shape the actual model code rejects. + return layer_idx >= first_shared > 0 + + +def _is_dense_mlp_layer(arch: ModelArchConfig, layer_idx: int) -> bool: + if arch.dense_layer_indices: + return layer_idx in arch.dense_layer_indices + return layer_idx < arch.num_dense_layers + + +def _per_layer_input_quantizable(arch: ModelArchConfig) -> int: + # why: Gemma4 PLE block adds per_layer_model_projection (single Linear), + # per_layer_input_gate (per layer), and per_layer_projection (per layer); + # see transformers gemma4/modular_gemma4.py:1077-1083 and :1247-1253. + pli = arch.hidden_size_per_layer_input + if pli <= 0: + return 0 + n_layers = arch.num_hidden_layers + hd = arch.hidden_size + return hd * (n_layers * pli) + (hd * pli) * n_layers + (pli * hd) * n_layers + + +def _per_layer_input_norm_elements(arch: ModelArchConfig) -> int: + pli = arch.hidden_size_per_layer_input + if pli <= 0: + return 0 + n_layers = arch.num_hidden_layers + hd = arch.hidden_size + return hd * n_layers + pli + + +def _per_layer_input_lora_params( + arch: ModelArchConfig, + r: int, + target_modules, +) -> int: + # why: Unsloth's get_peft_regex (unsloth_zoo/peft_utils.py) 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. Only count + # PLE LoRA when the user explicitly names PLE modules. + pli = arch.hidden_size_per_layer_input + if pli <= 0: + return 0 + targets = ( + {target_modules} + if isinstance(target_modules, str) + else set(target_modules or []) + ) + n_layers = arch.num_hidden_layers + hd = arch.hidden_size + total = 0 + if "per_layer_model_projection" in targets: + total += hd * r + r * (n_layers * pli) + if "per_layer_input_gate" in targets: + total += (hd * r + r * pli) * n_layers + if "per_layer_projection" in targets: + total += (pli * r + r * hd) * n_layers + return total + + +def _layer_attention_dims(arch: ModelArchConfig, layer_idx: int) -> tuple: + layer_types = _layer_types(arch) + layer_type = layer_types[layer_idx] + is_sliding = layer_type == "sliding_attention" + head_dim = ( + arch.global_head_dim + if not is_sliding and arch.global_head_dim + else _head_dim(arch) + ) + use_alt_attention = arch.attention_k_eq_v and not is_sliding + num_kv_heads = ( + arch.num_global_key_value_heads + if use_alt_attention and arch.num_global_key_value_heads + else arch.num_key_value_heads + ) + q_size = arch.num_attention_heads * head_dim + kv_size = num_kv_heads * head_dim + has_k = not _is_kv_shared_layer(arch, layer_idx) + has_v = has_k and not use_alt_attention + return q_size, kv_size, has_k, has_v + + +def _layer_mlp_size(arch: ModelArchConfig, layer_idx: int) -> int: + if arch.use_double_wide_mlp and _is_kv_shared_layer(arch, layer_idx): + return _dense_mlp_size(arch) * 2 + return _dense_mlp_size(arch) + + +def _text_linear_dims( + arch: ModelArchConfig, + layer_idx: int, +) -> Dict[str, tuple[int, int]]: + hd = arch.hidden_size + if _uses_structured_layer_shapes(arch): + q_size, kv_size, has_k, has_v = _layer_attention_dims(arch, layer_idx) + mlp_size = _layer_mlp_size(arch, layer_idx) + else: + q_size = hd + kv_size = _get_kv_size(arch) + has_k = True + has_v = True + mlp_size = _get_mlp_size(arch) + + dims = { + "q_proj": (hd, q_size), + "o_proj": (q_size, hd), + } + if has_k: + dims["k_proj"] = (hd, kv_size) + if has_v: + dims["v_proj"] = (hd, kv_size) + + dims.update( + { + "gate_proj": (hd, mlp_size), + "up_proj": (hd, mlp_size), + "down_proj": (mlp_size, hd), + } + ) + return dims + + +def _module_path_matches(skip_module: str, alias: str) -> bool: + skip_parts = [part for part in skip_module.split(".") if part] + alias_parts = [part for part in alias.split(".") if part] + if not skip_parts or not alias_parts: + return False + if alias_parts[0] == "layers": + return skip_parts == alias_parts + if len(skip_parts) <= len(alias_parts): + # why: transformers BNB quantizer suffix-matches short skip entries + # like ["q_proj"] / ["lm_head"] against full module paths, so a skip + # shorter than the alias is a tail match. + return alias_parts[-len(skip_parts) :] == skip_parts + if skip_parts[-len(alias_parts) :] != alias_parts: + return False + prefix_parts = skip_parts[: len(skip_parts) - len(alias_parts)] + if not prefix_parts: + return True + # why: bound the prefix to known text-tower roots so VLM skip names like + # vision_tower.model.layers..self_attn.q_proj do not shadow the text + # alias model.layers..self_attn.q_proj. + return ".".join(prefix_parts) in _SKIP_MODULE_TEXT_PREFIXES + + +def _add_module_aliases( + aliases: Dict[str, str], + canonical: str, + suffix: str, +) -> None: + for prefix in ( + "", + "model", + "model.model", + "language_model", + "language_model.model", + "model.language_model", + "model.language_model.model", + ): + alias = f"{prefix}.{suffix}" if prefix else suffix + aliases[alias] = canonical + + +def _build_text_module_elements( + arch: ModelArchConfig, +) -> tuple[Dict[str, int], Dict[str, str]]: + elements: Dict[str, int] = {} + aliases: Dict[str, str] = {} + + is_mla = arch.q_lora_rank is not None and not _uses_structured_layer_shapes(arch) + pli = arch.hidden_size_per_layer_input + hd_global = arch.hidden_size + + for layer_idx in range(arch.num_hidden_layers): + layer_modules: Dict[str, int] = {} + dims = _text_linear_dims(arch, layer_idx) + attn_dims = { + name: dim for name, dim in dims.items() if name in ATTENTION_TARGET_MODULES + } + mlp_dims = { + name: dim for name, dim in dims.items() if name in MLP_TARGET_MODULES + } + + if is_mla: + # why: _text_linear_dims uses (hd, hd) for q/o; MLA actually splits + # into q_a/q_b/kv_a/kv_b, so emit a single self_attn aggregate at + # the authoritative MLA per-layer total. + layer_modules["self_attn"] = _compute_attn_elements(arch) + else: + for name, (in_dim, out_dim) in attn_dims.items(): + layer_modules[f"self_attn.{name}"] = in_dim * out_dim + + if arch.num_experts and arch.num_experts > 1: + if _is_dense_mlp_layer(arch, layer_idx): + layer_modules.update( + { + f"mlp.{name}": in_dim * out_dim + for name, (in_dim, out_dim) in mlp_dims.items() + } + ) + else: + layer_modules["mlp.experts"] = _compute_routed_moe_elements(arch) + shared_moe = _compute_shared_moe_elements(arch) + if shared_moe: + # why: Qwen3.5-MoE exposes shared expert as + # mlp.shared_expert; Exaone-MoE/Laguna/GLM-style configs use + # mlp.shared_experts. Register both names so child-path + # llm_int8_skip_modules entries match the right shared block. + layer_modules["mlp.shared_expert"] = shared_moe + if arch.moe_has_dense_mlp: + # why: enable_moe_block runs the dense MLP and the MoE + # experts in parallel; register both for skip matching. + # Non-structured _text_linear_dims returns mlp_size from + # _get_mlp_size which prefers moe_intermediate_size, so + # rebuild dense dims from arch.intermediate_size directly. + if _uses_structured_layer_shapes(arch): + dense_dims = mlp_dims + else: + hd = arch.hidden_size + inter = arch.intermediate_size + dense_dims = { + "gate_proj": (hd, inter), + "up_proj": (hd, inter), + "down_proj": (inter, hd), + } + layer_modules.update( + { + f"mlp.{name}": in_dim * out_dim + for name, (in_dim, out_dim) in dense_dims.items() + } + ) + else: + layer_modules.update( + { + f"mlp.{name}": in_dim * out_dim + for name, (in_dim, out_dim) in mlp_dims.items() + } + ) + + if pli > 0: + # why: register PLE per-layer linears so llm_int8_skip_modules + # entries like model.layers.0.per_layer_input_gate match. + layer_modules["per_layer_input_gate"] = hd_global * pli + layer_modules["per_layer_projection"] = pli * hd_global + + attn_total = sum( + value + for name, value in layer_modules.items() + if name == "self_attn" or name.startswith("self_attn.") + ) + # why: gemma4 enable_moe_block puts routed experts at the sibling + # layers..experts attribute, not under self.mlp; the layer's "mlp" + # aggregate must reflect only the dense MLP path so a skip module + # `model.layers.0.mlp` does not over-skip into the experts block. + is_sibling_experts = bool(arch.moe_has_dense_mlp) + mlp_total = sum( + value + for name, value in layer_modules.items() + if ( + name == "mlp" + or ( + name.startswith("mlp.") + and not (is_sibling_experts and name == "mlp.experts") + ) + ) + ) + experts_total = layer_modules.get("mlp.experts", 0) if is_sibling_experts else 0 + layer_total = sum(layer_modules.values()) + + aggregate_modules = { + f"text.layers.{layer_idx}": layer_total, + f"text.layers.{layer_idx}.self_attn": attn_total, + f"text.layers.{layer_idx}.mlp": mlp_total, + } + if experts_total: + aggregate_modules[f"text.layers.{layer_idx}.experts"] = experts_total + elements.update(aggregate_modules) + for canonical in aggregate_modules: + suffix = canonical.removeprefix("text.") + _add_module_aliases(aliases, canonical, suffix) + + for name, value in layer_modules.items(): + canonical = f"text.layers.{layer_idx}.{name}" + elements[canonical] = value + _add_module_aliases(aliases, canonical, canonical.removeprefix("text.")) + if name == "mlp.experts" and arch.moe_has_dense_mlp: + # why: gemma4 enable_moe_block exposes routed experts at + # layers..experts (sibling of self.mlp), not under mlp. + _add_module_aliases(aliases, canonical, f"layers.{layer_idx}.experts") + elif name == "mlp.shared_expert": + # why: Exaone-MoE / Laguna / GLM-style configs use the plural + # `shared_experts` attribute name; register both spellings. + _add_module_aliases( + aliases, + canonical, + f"layers.{layer_idx}.mlp.shared_experts", + ) + + if pli > 0: + canonical = "text.per_layer_model_projection" + elements[canonical] = hd_global * (arch.num_hidden_layers * pli) + _add_module_aliases(aliases, canonical, canonical.removeprefix("text.")) + + return elements, aliases + + +def _compute_skipped_quantizable_elements(arch: ModelArchConfig) -> int: + if not arch.quantization_skip_modules: + return 0 + + module_elements, aliases = _build_text_module_elements(arch) + matched = set() + for skip_module in arch.quantization_skip_modules: + for alias, canonical in aliases.items(): + if _module_path_matches(skip_module, alias): + matched.add(canonical) + + pruned = { + canonical + for canonical in matched + if not any( + canonical != parent and canonical.startswith(f"{parent}.") + for parent in matched + ) + } + return sum(module_elements[canonical] for canonical in pruned) + + def _get_kv_size(arch: ModelArchConfig) -> int: return (arch.hidden_size // arch.num_attention_heads) * arch.num_key_value_heads @@ -226,6 +782,12 @@ def _get_mlp_size(arch: ModelArchConfig) -> int: return arch.intermediate_size +def _dense_mlp_size(arch: ModelArchConfig) -> int: + # why: Llama4 dense layers use intermediate_size_mlp; routed/shared + # experts use intermediate_size. Other configs leave the field None. + return arch.dense_intermediate_size or arch.intermediate_size + + def _get_num_experts(arch: ModelArchConfig) -> int: return arch.num_experts if arch.num_experts and arch.num_experts > 1 else 1 @@ -248,14 +810,39 @@ def _compute_attn_elements(arch: ModelArchConfig) -> int: def _compute_dense_mlp_elements(arch: ModelArchConfig) -> int: - return arch.hidden_size * arch.intermediate_size * 3 + return arch.hidden_size * _dense_mlp_size(arch) * 3 + + +def _shared_expert_size(arch: ModelArchConfig) -> int: + # why: Qwen3.5-MoE shared expert has its own intermediate_size (default 512) + # distinct from moe_intermediate_size; fall back to routed mlp_size for + # families that share it (deepseek-style configs). + return arch.shared_expert_intermediate_size or _get_mlp_size(arch) + + +def _compute_routed_moe_elements(arch: ModelArchConfig) -> int: + hd = arch.hidden_size + n_experts = _get_num_experts(arch) + return hd * _get_mlp_size(arch) * 3 * n_experts + n_experts * hd + + +def _compute_shared_moe_elements(arch: ModelArchConfig) -> int: + if not arch.n_shared_experts: + return 0 + hd = arch.hidden_size + shared_size = _shared_expert_size(arch) + total = hd * shared_size * 3 * arch.n_shared_experts + # why: only Qwen2-MoE / Qwen3.5-MoE define a shared_expert_gate Linear + # (hidden_size→1); other families (Exaone-MoE, HY-V3, GLM4-MoE-Lite, Laguna) + # have shared_experts without a gate. shared_expert_intermediate_size is the + # Qwen-style discriminator. + if arch.shared_expert_intermediate_size: + total += arch.n_shared_experts * hd + return total def _compute_moe_mlp_elements(arch: ModelArchConfig) -> int: - hd = arch.hidden_size - mlp_size = _get_mlp_size(arch) - n_experts = _get_num_experts(arch) - return hd * mlp_size * 3 * (n_experts + arch.n_shared_experts) + n_experts * hd + return _compute_routed_moe_elements(arch) + _compute_shared_moe_elements(arch) def _compute_layer_elements(arch: ModelArchConfig): @@ -267,22 +854,60 @@ def _compute_layer_elements(arch: ModelArchConfig): n_layers = arch.num_hidden_layers n_experts = _get_num_experts(arch) - attn_total = _compute_attn_elements(arch) * n_layers - - if n_experts > 1: + if _uses_structured_layer_shapes(arch): + attn_total = 0 + per_layer_dense_mlp = [] + for layer_idx in range(n_layers): + layer_dense_mlp = 0 + for name, (in_dim, out_dim) in _text_linear_dims( + arch, + layer_idx, + ).items(): + elements = in_dim * out_dim + if name in ATTENTION_TARGET_MODULES: + attn_total += elements + elif name in MLP_TARGET_MODULES: + layer_dense_mlp += elements + per_layer_dense_mlp.append(layer_dense_mlp) + if n_experts > 1: + n_dense = arch.num_dense_layers + n_moe = n_layers - n_dense + moe_mlp_total = _compute_moe_mlp_elements(arch) * n_moe + if arch.moe_has_dense_mlp: + # why: enable_moe_block runs dense MLP and MoE experts in + # parallel; count dense for every layer alongside MoE. + mlp_total = sum(per_layer_dense_mlp) + moe_mlp_total + else: + dense_only_total = sum( + value + for i, value in enumerate(per_layer_dense_mlp) + if _is_dense_mlp_layer(arch, i) + ) + mlp_total = moe_mlp_total + dense_only_total + else: + mlp_total = sum(per_layer_dense_mlp) + elif n_experts > 1: + attn_total = _compute_attn_elements(arch) * n_layers n_dense = arch.num_dense_layers n_moe = n_layers - n_dense - mlp_total = ( - _compute_moe_mlp_elements(arch) * n_moe - + _compute_dense_mlp_elements(arch) * n_dense - ) + moe_mlp_total = _compute_moe_mlp_elements(arch) * n_moe + if arch.moe_has_dense_mlp: + mlp_total = _compute_dense_mlp_elements(arch) * n_layers + moe_mlp_total + else: + mlp_total = moe_mlp_total + _compute_dense_mlp_elements(arch) * n_dense else: + attn_total = _compute_attn_elements(arch) * n_layers mlp_total = _compute_dense_mlp_elements(arch) * n_layers layernorms = 2 * hd - embed_tokens = arch.vocab_size * hd + per_layer_embed = ( + arch.vocab_size_per_layer_input * arch.hidden_size_per_layer_input * n_layers + ) + ple_text_linear = _per_layer_input_quantizable(arch) + ple_norms = _per_layer_input_norm_elements(arch) + embed_tokens = arch.vocab_size * hd + per_layer_embed + ple_norms lm_head = 0 if arch.tie_word_embeddings else arch.vocab_size * hd - return attn_total + mlp_total, layernorms, embed_tokens, lm_head + return attn_total + mlp_total + ple_text_linear, layernorms, embed_tokens, lm_head def compute_model_weights_bytes( @@ -295,7 +920,16 @@ def compute_model_weights_bytes( non_quantizable = layernorms * n_layers + embed_tokens + lm_head if training_method == "qlora" and load_in_4bit: - return int(total_quantizable * 2 / QUANT_4BIT_FACTOR + non_quantizable * 2) + skipped_quantizable = min( + _compute_skipped_quantizable_elements(arch), + total_quantizable, + ) + quantized = total_quantizable - skipped_quantizable + return int( + quantized * 2 / arch.quant_4bit_factor + + skipped_quantizable * 2 + + non_quantizable * 2 + ) return int((total_quantizable + non_quantizable) * 2) @@ -363,46 +997,130 @@ def compute_lora_params( lora_rank: int, target_modules: list, ) -> int: + all_linear = _targets_all_linear(target_modules) + selected_modules = list(DEFAULT_TARGET_MODULES) if all_linear else target_modules hd = arch.hidden_size r = lora_rank n_layers = arch.num_hidden_layers n_experts = _get_num_experts(arch) - attn_total = _lora_attn_elements(arch, r, target_modules) * n_layers - - if n_experts > 1: + use_structured_shapes = _uses_structured_layer_shapes(arch) + if use_structured_shapes: + attn_total = 0 + structured_dense_mlp = 0 + per_layer_dense_mlp = [] + for layer_idx in range(n_layers): + layer_dense = 0 + for name, (in_dim, out_dim) in _text_linear_dims( + arch, + layer_idx, + ).items(): + if name not in selected_modules: + continue + if name in ATTENTION_TARGET_MODULES: + attn_total += in_dim * r + r * out_dim + elif name in MLP_TARGET_MODULES: + layer_dense += in_dim * r + r * out_dim + per_layer_dense_mlp.append(layer_dense) + structured_dense_mlp += layer_dense + if n_experts > 1: + n_dense = arch.num_dense_layers + n_moe = n_layers - n_dense + # why: peft "all-linear" attaches LoRA to nn.Linear only; + # routed experts are nn.Parameter and need explicit + # gate_proj/up_proj/down_proj naming via Unsloth's + # get_moe_target_parameters. Shared experts are nn.Linear and + # are picked up by get_peft_regex. + routed_moe = ( + 0 + if all_linear + else _lora_mlp_elements( + hd, + _get_mlp_size(arch), + r, + selected_modules, + n_experts, + ) + ) + shared_moe = _lora_mlp_elements( + hd, + _shared_expert_size(arch), + r, + selected_modules, + arch.n_shared_experts, + ) + moe_mlp = routed_moe + shared_moe + if arch.moe_has_dense_mlp: + # why: parallel dense MLP coexists with MoE on every layer. + mlp_total = structured_dense_mlp + moe_mlp * n_moe + else: + dense_only = sum( + value + for i, value in enumerate(per_layer_dense_mlp) + if _is_dense_mlp_layer(arch, i) + ) + mlp_total = moe_mlp * n_moe + dense_only + else: + mlp_total = structured_dense_mlp + return ( + attn_total + + mlp_total + + _per_layer_input_lora_params(arch, r, target_modules) + ) + elif n_experts > 1: + attn_total = _lora_attn_elements(arch, r, selected_modules) * n_layers n_dense = arch.num_dense_layers n_moe = n_layers - n_dense - # Include shared experts alongside routed experts - moe_expert_mult = n_experts + arch.n_shared_experts - moe_mlp = _lora_mlp_elements( - hd, - _get_mlp_size(arch), - r, - target_modules, - moe_expert_mult, + # why: routed and shared experts may use different intermediate sizes + # (Qwen3.5-MoE: routed mlp_size != shared_expert_intermediate_size). + # See structured branch for the all-linear exclusion rationale; only + # routed (nn.Parameter) experts are excluded under all-linear. + routed_moe = ( + 0 + if all_linear + else _lora_mlp_elements( + hd, + _get_mlp_size(arch), + r, + selected_modules, + n_experts, + ) ) + shared_moe = _lora_mlp_elements( + hd, + _shared_expert_size(arch), + r, + selected_modules, + arch.n_shared_experts, + ) + moe_mlp = routed_moe + shared_moe dense_mlp = _lora_mlp_elements( hd, - arch.intermediate_size, + _dense_mlp_size(arch), r, - target_modules, + selected_modules, 1, ) - mlp_total = moe_mlp * n_moe + dense_mlp * n_dense + if arch.moe_has_dense_mlp: + mlp_total = moe_mlp * n_moe + dense_mlp * n_layers + else: + mlp_total = moe_mlp * n_moe + dense_mlp * n_dense else: + attn_total = _lora_attn_elements(arch, r, selected_modules) * n_layers mlp_total = ( _lora_mlp_elements( hd, - arch.intermediate_size, + _dense_mlp_size(arch), r, - target_modules, + selected_modules, 1, ) * n_layers ) - return attn_total + mlp_total + return ( + attn_total + mlp_total + _per_layer_input_lora_params(arch, r, target_modules) + ) def compute_lora_adapter_bytes(lora_params: int) -> int: @@ -419,26 +1137,88 @@ def compute_gradient_bytes(trainable_params: int) -> int: return trainable_params * 2 +def _is_linear_attention(attention_implementation: Optional[str]) -> bool: + # why: PyTorch SDPA dispatches to flash/memory-efficient O(n) backends; only + # eager (and other non-flash impls) need the quadratic correction. + return attention_implementation in LINEAR_ATTENTION_IMPLS + + +def _compute_non_flash_attention_bytes( + arch: ModelArchConfig, + batch_size: int, + seq_len: int, + effective_layers: float, +) -> int: + score_elements = batch_size * arch.num_attention_heads * seq_len * seq_len + return int(score_elements * 2 * NON_FLASH_ATTENTION_FACTOR * effective_layers) + + +def _layer_qkv_mlp_sizes(arch: ModelArchConfig, layer_idx: int) -> tuple: + n_experts = _get_num_experts(arch) + is_moe_layer = n_experts > 1 and not _is_dense_mlp_layer(arch, layer_idx) + if _uses_structured_layer_shapes(arch): + q_size, kv_size, _has_k, _has_v = _layer_attention_dims(arch, layer_idx) + # why: KV-shared layers (Gemma4/Gemma3n) drop k_proj/v_proj WEIGHTS but + # the donor layer's K/V tensors stay alive across the shared range, so + # activation memory still pays for kv_size; only the weight path uses + # has_k/has_v. + layer_type = _layer_types(arch)[layer_idx] + use_alt_attention = arch.attention_k_eq_v and layer_type != "sliding_attention" + kv_count = 1 if use_alt_attention else 2 + qkv_size = q_size + kv_size * kv_count + if is_moe_layer: + # why: each token routes through `num_experts_per_tok` experts; their + # gate/up/down intermediates are all live during MLP forward. + mlp_size = _get_mlp_size(arch) * arch.num_experts_per_tok + if arch.n_shared_experts: + mlp_size += _shared_expert_size(arch) * arch.n_shared_experts + if arch.moe_has_dense_mlp: + mlp_size += _layer_mlp_size(arch, layer_idx) + else: + mlp_size = _layer_mlp_size(arch, layer_idx) + return qkv_size, mlp_size + kv_size = _get_kv_size(arch) + if is_moe_layer: + mlp_size = _get_mlp_size(arch) * arch.num_experts_per_tok + if arch.n_shared_experts: + mlp_size += _shared_expert_size(arch) * arch.n_shared_experts + if arch.moe_has_dense_mlp: + mlp_size += arch.intermediate_size + else: + mlp_size = _get_mlp_size(arch) + return arch.hidden_size + kv_size + kv_size, mlp_size + + +def _per_layer_activation_bytes( + arch: ModelArchConfig, + layer_idx: int, + batch_size: int, + seq_len: int, +) -> int: + qkv_size, mlp_size = _layer_qkv_mlp_sizes(arch, layer_idx) + activation_qkv = seq_len * batch_size * qkv_size + residual_memory = (seq_len * batch_size) * 2 + activation_mlp = seq_len * batch_size * (mlp_size + mlp_size) + # why: per_layer_input_gate (hd-sized) and per_layer_projection (pli-sized) + # outputs materialize once per decoder layer when hidden_size_per_layer_input + # is set; see gemma4/modular_gemma4.py:1141-1145. + pli = arch.hidden_size_per_layer_input + activation_ple = seq_len * batch_size * (arch.hidden_size + pli) if pli > 0 else 0 + return int( + (activation_qkv + residual_memory + activation_mlp + activation_ple) * 2 * 1.25 + ) + + def compute_activation_bytes( arch: ModelArchConfig, batch_size: int, seq_len: int, gradient_checkpointing: str, is_lora: bool = False, + attention_implementation: Optional[str] = "flash_attention_2", ) -> int: - hd = arch.hidden_size - kv_size = _get_kv_size(arch) - mlp_size = _get_mlp_size(arch) - bsz = batch_size n_layers = arch.num_hidden_layers - activation_qkv = seq_len * bsz * (hd + kv_size + kv_size) - residual_memory = (seq_len * bsz) * 2 - activation_mlp = seq_len * bsz * (mlp_size + mlp_size) - - per_layer_bytes = (activation_qkv + residual_memory + activation_mlp) * 2 - per_layer_bytes = int(per_layer_bytes * 1.25) - gc_key = gradient_checkpointing.lower() gc_entry = GC_LAYER_MULTIPLIERS.get(gc_key, (None, None)) full_ft_mult, lora_mult = gc_entry @@ -446,10 +1226,35 @@ def compute_activation_bytes( if gc_multiplier is None: effective_layers = n_layers + linear_bytes = sum( + _per_layer_activation_bytes(arch, i, batch_size, seq_len) + for i in range(n_layers) + ) else: effective_layers = gc_multiplier + max_layer_bytes = max( + _per_layer_activation_bytes(arch, i, batch_size, seq_len) + for i in range(n_layers) + ) + linear_bytes = int(max_layer_bytes * effective_layers) - return int(per_layer_bytes * effective_layers) + # why: gemma4 per_layer_model_projection runs once outside the per-decoder + # loop and materializes a [B, S, L, PLI] tensor; see modular_gemma4.py:1247. + pli = arch.hidden_size_per_layer_input + if pli > 0: + linear_bytes += int(seq_len * batch_size * n_layers * pli * 2 * 1.25) + + if _is_linear_attention(attention_implementation): + return linear_bytes + return max( + linear_bytes, + _compute_non_flash_attention_bytes( + arch, + batch_size, + seq_len, + effective_layers, + ), + ) def estimate_training_vram( @@ -474,21 +1279,23 @@ def estimate_training_vram( trainable_params = lora_params if is_lora else compute_total_params(arch) optimizer_bytes = compute_optimizer_bytes(trainable_params, config.optimizer) - gradient_bytes = max( - compute_gradient_bytes(trainable_params), - int(model_weights * 0.15), - ) activations_computed = compute_activation_bytes( arch, config.batch_size, config.max_seq_length, config.gradient_checkpointing, is_lora = is_lora, + attention_implementation = config.attention_implementation, ) - activation_bytes = max( - activations_computed, - int(model_weights * 0.15 * (config.batch_size / 2)), - ) + raw_gradient_bytes = compute_gradient_bytes(trainable_params) + gradient_floor = int(model_weights * 0.15) + if is_lora: + gradient_floor = min( + gradient_floor, + max(activations_computed, optimizer_bytes), + ) + gradient_bytes = max(raw_gradient_bytes, gradient_floor) + activation_bytes = activations_computed return VramBreakdown( model_weights = model_weights,