Views
No views yet
0.16.0.dev01base_model: Qwen/Qwen3-30B-A3B
2
3# GKD plugin: distill from full (unpruned) teacher into pruned student
4plugins:
5 - axolotl.integrations.gkd.GKDPlugin
6 # - axolotl.integrations.kernels.KernelsPlugin # for ScatterMoE kernels (optional)
7
8strict: false
9
10# GKD trainer config
11gkd_trainer: true
12gkd_teacher_model: Qwen/Qwen3.6-35B-A3B
13gkd_temperature: 1.0 # T=1.0 optimal for LLMs (Minitron); high-entropy logits don't need softening
14gkd_distill_alpha: 1.0 # weight for distillation KL loss
15gkd_ce_alpha: 0.1 # pure logit distillation (Minitron: L_logits alone > L_CLM + L_logits for pruned recovery)
16gkd_scale_kl_with_temperature_sq: true # standard KD gradient scaling (T^2); no-op at T=1
17
18chat_template: qwen3_5
19train_on_inputs: true
20datasets:
21 - path: /data/workspace/kunato/quantize/reap2/artifacts/data/calibration/merged-calibration.jsonl
22 ds_type: json
23 type: chat_template
24 split: train
25 field_messages: messages
26 split_thinking: true
27val_set_size: 0.0
28output_dir: ./outputs/qwen3.6-35b-a3b-pruned-gkd-typhoon-math-prune-v2
29dataset_prepared_path: last_run_prepared
30
31sequence_len: 16384
32sample_packing: true
33
34load_in_4bit: true
35quantize_moe_experts: true
36adapter: qlora
37lora_r: 32
38lora_alpha: 64
39lora_dropout: 0
40
41# Target routed expert 3D params via lora_target_parameters (requires lora_dropout: 0)
42lora_target_parameters:
43 - mlp.experts.gate_up_proj
44 - mlp.experts.down_proj
45
46# # Target shared expert standard Linear4bit layers via lora_target_modules
47# # (lora_target_parameters can't match standard nn.Linear modules)
48# lora_target_modules:
49# - shared_expert.gate_proj
50# - shared_expert.up_proj
51# - shared_expert.down_proj
52
53lora_target_modules:
54 - q_proj
55 - k_proj
56 - v_proj
57 - o_proj
58
59wandb_project:
60wandb_entity:
61wandb_watch:
62wandb_name:
63wandb_log_model:
64
65gradient_accumulation_steps: 1
66micro_batch_size: 1
67num_epochs: 3
68optimizer: adamw_torch
69lr_scheduler: cosine
70learning_rate: 0.0002
71
72bf16: auto
73tf32: true
74
75# ScatterMoE: fused Triton kernels for MoE expert forward/backward with LoRA.
76# Using these kernels doesn't improve the training speed in this setting.
77# Bypasses ParamWrapper, computing Y = X@W + scaling*(X@A^T)@B^T in one pass.
78# use_scattermoe: true
79# Only dequantize routed experts (e.g. 8 of 256), ~97% memory savings per layer.
80# selective_expert_dequant: true
81
82lora_mlp_kernel: false
83lora_qkv_kernel: false
84lora_o_kernel: false
85
86# Use FSDP activation_checkpointing (not gradient_checkpointing) to avoid
87# the HF ValueError when both are true simultaneously.
88gradient_checkpointing: false
89gradient_checkpointing_kwargs:
90 use_reentrant: false
91resume_from_checkpoint:
92logging_steps: 1
93flash_attention: true
94
95warmup_ratio: 0.1
96evals_per_epoch: 0
97saves_per_epoch: 1
98save_total_limit: 1
99weight_decay: 0.0
100special_tokens:
101
102fsdp_config:
103 fsdp_version: 2
104 offload_params: true
105 cpu_ram_efficient_loading: false
106 auto_wrap_policy: TRANSFORMER_BASED_WRAP
107 transformer_layer_cls_to_wrap: Qwen3_5MoeDecoderLayer
108 state_dict_type: FULL_STATE_DICT
109 sharding_strategy: FULL_SHARD
110 reshard_after_forward: true
111 activation_checkpointing: true
112