# SPDX-License-Identifier: AGPL-3.0-only # 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, TrainingVramConfig, extract_arch_config, compute_model_weights_bytes, compute_total_params, compute_lora_params, compute_lora_adapter_bytes, compute_optimizer_bytes, compute_gradient_bytes, compute_activation_bytes, estimate_training_vram, DEFAULT_TARGET_MODULES, ) def _gb(b: int) -> float: return b / (1024**3) LLAMA_8B = ModelArchConfig( 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, ) QWEN_05B = ModelArchConfig( hidden_size = 896, num_hidden_layers = 24, num_attention_heads = 14, num_key_value_heads = 2, intermediate_size = 4864, vocab_size = 151936, tie_word_embeddings = True, ) MOE_CONFIG = ModelArchConfig( hidden_size = 4096, num_hidden_layers = 32, num_attention_heads = 32, num_key_value_heads = 8, intermediate_size = 14336, vocab_size = 32000, tie_word_embeddings = False, num_experts = 8, ) DEEPSEEK_V3 = ModelArchConfig( hidden_size = 7168, num_hidden_layers = 61, num_attention_heads = 128, num_key_value_heads = 128, intermediate_size = 18432, vocab_size = 129280, tie_word_embeddings = False, num_experts = 256, moe_intermediate_size = 2048, n_shared_experts = 1, num_dense_layers = 3, q_lora_rank = 1536, kv_lora_rank = 512, qk_nope_head_dim = 128, qk_rope_head_dim = 64, v_head_dim = 128, ) QWEN3_MOE_30B = ModelArchConfig( hidden_size = 2048, num_hidden_layers = 48, num_attention_heads = 32, num_key_value_heads = 4, intermediate_size = 8192, vocab_size = 151936, tie_word_embeddings = True, num_experts = 128, moe_intermediate_size = 768, n_shared_experts = 0, num_dense_layers = 0, ) GLM4_MOE = ModelArchConfig( hidden_size = 4096, num_hidden_layers = 46, num_attention_heads = 96, num_key_value_heads = 8, intermediate_size = 10944, vocab_size = 151552, tie_word_embeddings = False, num_experts = 128, moe_intermediate_size = 1408, n_shared_experts = 1, num_dense_layers = 1, ) GPT_OSS = ModelArchConfig( hidden_size = 6144, num_hidden_layers = 64, num_attention_heads = 64, num_key_value_heads = 8, intermediate_size = 2880, vocab_size = 200064, tie_word_embeddings = False, num_experts = 128, moe_intermediate_size = None, n_shared_experts = 0, 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): hf_config = 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, ) arch = extract_arch_config(hf_config) self.assertIsNotNone(arch) self.assertEqual(arch.hidden_size, 4096) self.assertEqual(arch.num_hidden_layers, 32) self.assertEqual(arch.num_key_value_heads, 8) self.assertIsNone(arch.num_experts) def test_vlm_text_config(self): text_cfg = SimpleNamespace( hidden_size = 2048, num_hidden_layers = 24, num_attention_heads = 16, num_key_value_heads = 4, intermediate_size = 8192, vocab_size = 32000, tie_word_embeddings = True, ) hf_config = SimpleNamespace(text_config = text_cfg) arch = extract_arch_config(hf_config) self.assertIsNotNone(arch) self.assertEqual(arch.hidden_size, 2048) def test_moe_detection(self): hf_config = SimpleNamespace( hidden_size = 4096, num_hidden_layers = 32, num_attention_heads = 32, num_key_value_heads = 8, intermediate_size = 14336, vocab_size = 32000, tie_word_embeddings = False, num_local_experts = 8, ) arch = extract_arch_config(hf_config) self.assertEqual(arch.num_experts, 8) def test_missing_fields_returns_none(self): hf_config = SimpleNamespace(hidden_size = 4096) arch = extract_arch_config(hf_config) self.assertIsNone(arch) def test_intermediate_size_list(self): hf_config = SimpleNamespace( hidden_size = 2048, num_hidden_layers = 24, num_attention_heads = 16, num_key_value_heads = 4, intermediate_size = [8192, 8192], vocab_size = 32000, tie_word_embeddings = True, ) 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): weight_bytes = compute_model_weights_bytes(LLAMA_8B, "full", False) weight_gb = _gb(weight_bytes) self.assertGreater(weight_gb, 14.0) self.assertLess(weight_gb, 18.0) def test_llama_8b_qlora_4bit(self): weight_bytes = compute_model_weights_bytes(LLAMA_8B, "qlora", True) weight_gb = _gb(weight_bytes) self.assertGreater(weight_gb, 4.0) self.assertLess(weight_gb, 7.0) def test_4bit_smaller_than_fp16(self): fp16 = compute_model_weights_bytes(LLAMA_8B, "full", False) q4 = compute_model_weights_bytes(LLAMA_8B, "qlora", True) self.assertLess(q4, fp16) ratio = fp16 / q4 self.assertGreater(ratio, 2.0) self.assertLess(ratio, 4.0) def test_moe_larger_than_dense(self): dense = compute_model_weights_bytes(LLAMA_8B, "full", False) moe = compute_model_weights_bytes(MOE_CONFIG, "full", False) self.assertGreater(moe, dense * 3) class TestLoraParams(unittest.TestCase): def test_llama_8b_default_modules_rank16(self): lora_p = compute_lora_params(LLAMA_8B, 16, DEFAULT_TARGET_MODULES) total_p = compute_total_params(LLAMA_8B) ratio = lora_p / total_p self.assertGreater(ratio, 0.005) self.assertLess(ratio, 0.05) def test_higher_rank_more_params(self): r16 = compute_lora_params(LLAMA_8B, 16, DEFAULT_TARGET_MODULES) r64 = compute_lora_params(LLAMA_8B, 64, DEFAULT_TARGET_MODULES) self.assertAlmostEqual(r64 / r16, 4.0, places = 1) def test_fewer_modules_fewer_params(self): all_mods = compute_lora_params(LLAMA_8B, 16, DEFAULT_TARGET_MODULES) qv_only = compute_lora_params(LLAMA_8B, 16, ["q_proj", "v_proj"]) self.assertLess(qv_only, all_mods) def test_moe_mlp_modules_scale_with_experts(self): dense_lora = compute_lora_params(LLAMA_8B, 16, ["gate_proj", "up_proj", "down_proj"]) moe_lora = compute_lora_params(MOE_CONFIG, 16, ["gate_proj", "up_proj", "down_proj"]) 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"]) moe_attn = compute_lora_params(MOE_CONFIG, 16, ["q_proj", "k_proj", "v_proj", "o_proj"]) 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): self.assertEqual(compute_optimizer_bytes(1_000_000, "adamw_8bit"), 4_000_000) def test_adamw_torch(self): self.assertEqual(compute_optimizer_bytes(1_000_000, "adamw_torch"), 6_000_000) def test_sgd(self): self.assertEqual(compute_optimizer_bytes(1_000_000, "sgd"), 4_000_000) def test_unknown_defaults_to_4(self): self.assertEqual(compute_optimizer_bytes(1_000_000, "some_new_opt"), 4_000_000) class TestGradientBytes(unittest.TestCase): def test_fp16_gradients(self): self.assertEqual(compute_gradient_bytes(1_000_000), 2_000_000) class TestActivationBytes(unittest.TestCase): def test_no_gc_scales_with_layers(self): act_none = compute_activation_bytes(LLAMA_8B, 2, 2048, "none") act_gc = compute_activation_bytes(LLAMA_8B, 2, 2048, "true") self.assertGreater(act_none, act_gc * 10) def test_unsloth_gc_smaller_than_standard(self): act_true = compute_activation_bytes(LLAMA_8B, 2, 2048, "true") act_unsloth = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth") self.assertLess(act_unsloth, act_true) def test_lora_activations_smaller_than_full_ft(self): full_ft = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth", is_lora = False) lora = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth", is_lora = True) self.assertLess(lora, full_ft) def test_scales_with_batch_size(self): act_bsz2 = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth") act_bsz4 = compute_activation_bytes(LLAMA_8B, 4, 2048, "unsloth") self.assertAlmostEqual(act_bsz4 / act_bsz2, 2.0, delta = 0.1) def test_scales_with_seq_len(self): act_2k = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth") 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 with 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): config = TrainingVramConfig( training_method = "qlora", batch_size = 2, max_seq_length = 2048, lora_rank = 16, gradient_checkpointing = "unsloth", optimizer = "adamw_8bit", load_in_4bit = True, ) breakdown = estimate_training_vram(LLAMA_8B, config) total_gb = _gb(breakdown.total) self.assertGreater(total_gb, 5.0) self.assertLess(total_gb, 12.0) def test_llama_8b_full_ft_reasonable_total(self): config = TrainingVramConfig( training_method = "full", batch_size = 2, max_seq_length = 2048, gradient_checkpointing = "unsloth", optimizer = "adamw_8bit", load_in_4bit = False, ) breakdown = estimate_training_vram(LLAMA_8B, config) total_gb = _gb(breakdown.total) self.assertGreater(total_gb, 50.0) self.assertLess(total_gb, 75.0) def test_qlora_much_less_than_full_ft(self): qlora_config = TrainingVramConfig( training_method = "qlora", load_in_4bit = True, batch_size = 2, max_seq_length = 2048, ) full_config = TrainingVramConfig( training_method = "full", load_in_4bit = False, batch_size = 2, max_seq_length = 2048, ) qlora = estimate_training_vram(LLAMA_8B, qlora_config) full = estimate_training_vram(LLAMA_8B, full_config) self.assertLess(qlora.total, full.total / 3) def test_qwen_05b_qlora_fits_in_4gb(self): config = TrainingVramConfig( training_method = "qlora", batch_size = 2, max_seq_length = 2048, lora_rank = 16, gradient_checkpointing = "unsloth", optimizer = "adamw_8bit", load_in_4bit = True, ) breakdown = estimate_training_vram(QWEN_05B, config) total_gb = _gb(breakdown.total) self.assertLess(total_gb, 5.0) def test_breakdown_components_positive(self): config = TrainingVramConfig(training_method = "qlora", load_in_4bit = True) breakdown = estimate_training_vram(LLAMA_8B, config) self.assertGreater(breakdown.model_weights, 0) self.assertGreater(breakdown.lora_adapters, 0) self.assertGreater(breakdown.optimizer_states, 0) self.assertGreater(breakdown.gradients, 0) self.assertGreater(breakdown.activations, 0) self.assertGreater(breakdown.cuda_overhead, 0) def test_full_ft_no_lora_adapters(self): config = TrainingVramConfig(training_method = "full", load_in_4bit = False) breakdown = estimate_training_vram(LLAMA_8B, config) self.assertEqual(breakdown.lora_adapters, 0) def test_to_gb_dict_keys(self): config = TrainingVramConfig(training_method = "qlora", load_in_4bit = True) breakdown = estimate_training_vram(LLAMA_8B, config) gb_dict = breakdown.to_gb_dict() expected_keys = { "model_weights_gb", "lora_adapters_gb", "optimizer_states_gb", "gradients_gb", "activations_gb", "cuda_overhead_gb", "total_gb", } self.assertEqual(set(gb_dict.keys()), expected_keys) def test_total_equals_sum_of_parts(self): config = TrainingVramConfig(training_method = "qlora", load_in_4bit = True) breakdown = estimate_training_vram(LLAMA_8B, config) parts_sum = ( breakdown.model_weights + breakdown.lora_adapters + breakdown.optimizer_states + breakdown.gradients + breakdown.activations + breakdown.cuda_overhead ) self.assertEqual(breakdown.total, parts_sum) def test_larger_batch_increases_total(self): small = TrainingVramConfig( training_method = "qlora", load_in_4bit = True, batch_size = 1, ) large = TrainingVramConfig( training_method = "qlora", load_in_4bit = True, batch_size = 8, ) small_v = estimate_training_vram(LLAMA_8B, small) large_v = estimate_training_vram(LLAMA_8B, large) self.assertGreater(large_v.total, small_v.total) def test_adamw_fp32_uses_more_optimizer_memory(self): opt8 = TrainingVramConfig( training_method = "full", load_in_4bit = False, optimizer = "adamw_8bit", ) opt32 = TrainingVramConfig( training_method = "full", load_in_4bit = False, optimizer = "adamw_torch", ) v8 = estimate_training_vram(LLAMA_8B, opt8) v32 = estimate_training_vram(LLAMA_8B, opt32) self.assertAlmostEqual(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): hf_config = SimpleNamespace( hidden_size = 7168, num_hidden_layers = 61, num_attention_heads = 128, num_key_value_heads = 128, intermediate_size = 18432, vocab_size = 129280, tie_word_embeddings = False, n_routed_experts = 256, moe_intermediate_size = 2048, n_shared_experts = 1, first_k_dense_replace = 3, q_lora_rank = 1536, kv_lora_rank = 512, qk_nope_head_dim = 128, qk_rope_head_dim = 64, v_head_dim = 128, ) arch = extract_arch_config(hf_config) self.assertEqual(arch.num_experts, 256) self.assertEqual(arch.n_shared_experts, 1) self.assertEqual(arch.num_dense_layers, 3) self.assertEqual(arch.q_lora_rank, 1536) self.assertEqual(arch.kv_lora_rank, 512) def test_qwen3_moe_decoder_sparse_step(self): hf_config = SimpleNamespace( hidden_size = 2048, num_hidden_layers = 48, num_attention_heads = 32, num_key_value_heads = 4, intermediate_size = 8192, vocab_size = 151936, tie_word_embeddings = True, num_local_experts = 128, 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( hidden_size = 2048, num_hidden_layers = 24, num_attention_heads = 16, num_key_value_heads = 4, intermediate_size = 8192, vocab_size = 151936, tie_word_embeddings = True, num_local_experts = 60, moe_intermediate_size = 1408, decoder_sparse_step = 1, mlp_only_layers = [0, 1, 2, 3], ) arch = extract_arch_config(hf_config) self.assertEqual(arch.num_dense_layers, 4) def test_glm4_moe_first_k_dense(self): hf_config = SimpleNamespace( hidden_size = 4096, num_hidden_layers = 46, num_attention_heads = 96, num_key_value_heads = 8, intermediate_size = 10944, vocab_size = 151552, tie_word_embeddings = False, n_routed_experts = 128, moe_intermediate_size = 1408, n_shared_experts = 1, first_k_dense_replace = 1, ) arch = extract_arch_config(hf_config) self.assertEqual(arch.num_dense_layers, 1) self.assertEqual(arch.n_shared_experts, 1) def test_gpt_oss_no_moe_intermediate(self): hf_config = SimpleNamespace( hidden_size = 6144, num_hidden_layers = 64, num_attention_heads = 64, num_key_value_heads = 8, intermediate_size = 2880, vocab_size = 200064, tie_word_embeddings = False, num_local_experts = 128, ) arch = extract_arch_config(hf_config) self.assertEqual(arch.num_experts, 128) self.assertIsNone(arch.moe_intermediate_size) self.assertEqual(arch.num_dense_layers, 0) def test_backward_compat_no_new_fields(self): hf_config = 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, ) arch = extract_arch_config(hf_config) 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, `self.experts` is a sibling of # `self.mlp`; the `text.layers..mlp` aggregate covers 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) # PLE block adds these quantizable text linears: model_projection # (hd*nl*pli), per_layer_input_gate (hd*pli per layer), # per_layer_projection (pli*hd per layer). 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 a component tag (mlp/attn/...) # in module names; PLE names (per_layer_input_gate, per_layer_projection, # per_layer_model_projection) lack one, so all-linear 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` covers 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): def test_shared_experts_increase_weight_bytes(self): no_shared = ModelArchConfig( hidden_size = 4096, num_hidden_layers = 32, num_attention_heads = 32, num_key_value_heads = 8, intermediate_size = 14336, vocab_size = 32000, tie_word_embeddings = False, num_experts = 64, moe_intermediate_size = 1407, n_shared_experts = 0, ) with_shared = ModelArchConfig( hidden_size = 4096, num_hidden_layers = 32, num_attention_heads = 32, num_key_value_heads = 8, intermediate_size = 14336, vocab_size = 32000, tie_word_embeddings = False, num_experts = 64, moe_intermediate_size = 1407, n_shared_experts = 2, ) w_no = compute_model_weights_bytes(no_shared, "full", False) w_yes = compute_model_weights_bytes(with_shared, "full", False) self.assertGreater(w_yes, w_no) delta_per_layer = 4096 * 1407 * 3 * 2 expected_delta = delta_per_layer * 32 * 2 actual_delta = w_yes - w_no self.assertAlmostEqual(actual_delta, expected_delta, delta = expected_delta * 0.01) def test_deepseek_v3_params_in_range(self): total = compute_total_params(DEEPSEEK_V3) total_b = total / 1e9 self.assertGreater(total_b, 600) self.assertLess(total_b, 750) class TestMLA(unittest.TestCase): def test_mla_different_from_standard(self): from utils.hardware.vram_estimation import _compute_attn_elements mla_arch = DEEPSEEK_V3 std_arch = ModelArchConfig( hidden_size = 7168, num_hidden_layers = 61, num_attention_heads = 128, num_key_value_heads = 128, intermediate_size = 18432, vocab_size = 129280, ) mla_attn = _compute_attn_elements(mla_arch) std_attn = _compute_attn_elements(std_arch) self.assertNotEqual(mla_attn, std_attn) def test_mla_lora_produces_values(self): 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): all_moe = ModelArchConfig( hidden_size = 4096, num_hidden_layers = 46, num_attention_heads = 96, num_key_value_heads = 8, intermediate_size = 10944, vocab_size = 151552, tie_word_embeddings = False, num_experts = 128, moe_intermediate_size = 1408, n_shared_experts = 1, num_dense_layers = 0, ) mixed = ModelArchConfig( hidden_size = 4096, num_hidden_layers = 46, num_attention_heads = 96, num_key_value_heads = 8, intermediate_size = 10944, vocab_size = 151552, tie_word_embeddings = False, num_experts = 128, moe_intermediate_size = 1408, n_shared_experts = 1, num_dense_layers = 1, ) w_all = compute_model_weights_bytes(all_moe, "full", False) w_mixed = compute_model_weights_bytes(mixed, "full", False) self.assertNotEqual(w_all, w_mixed) def test_glm4_moe_params_reasonable(self): total = compute_total_params(GLM4_MOE) total_b = total / 1e9 self.assertGreater(total_b, 80) self.assertLess(total_b, 120) def test_qwen3_moe_30b_params_reasonable(self): total = compute_total_params(QWEN3_MOE_30B) total_b = total / 1e9 self.assertGreater(total_b, 20) self.assertLess(total_b, 50) def test_gpt_oss_uses_intermediate_size(self): total = compute_total_params(GPT_OSS) total_b = total / 1e9 self.assertGreater(total_b, 350) self.assertLess(total_b, 500) def test_lora_dense_vs_moe_layers_differ(self): all_moe = ModelArchConfig( hidden_size = 4096, num_hidden_layers = 10, num_attention_heads = 32, num_key_value_heads = 8, intermediate_size = 14336, vocab_size = 32000, tie_word_embeddings = False, num_experts = 8, moe_intermediate_size = 1024, num_dense_layers = 0, ) mixed = ModelArchConfig( hidden_size = 4096, num_hidden_layers = 10, num_attention_heads = 32, num_key_value_heads = 8, intermediate_size = 14336, vocab_size = 32000, tie_word_embeddings = False, num_experts = 8, moe_intermediate_size = 1024, num_dense_layers = 5, ) lora_all = compute_lora_params(all_moe, 16, ["gate_proj", "up_proj", "down_proj"]) lora_mix = compute_lora_params(mixed, 16, ["gate_proj", "up_proj", "down_proj"]) 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"])) # text.layers.0 aggregate includes the PLE per-layer modules, so the # same skip on a no-PLE config 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 # give 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 -> 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 a hypothetical shared expert: dense + # routed + shared all live per layer; mlp_size sums all three. 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 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]; element 1 # is the vision-routed width, not the shared-expert width. Shared # experts size 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()