Implement 4-bit quant for train script and inference

tools/train/train/qwq.py

@@ -30,7 +30,7 @@ class Args:
         parser.add_argument(
             '--base-model',
             type=str,
-            default='Qwen/QwQ-32B',
+            default='unsloth/QwQ-32B-bnb-4bit',
             help='HuggingFace model ID for base model'
         )
         parser.add_argument(
@@ -71,9 +71,6 @@ def main():
     dataset = Dataset(args.config_path)
     dataset.validate()
 
-    max_seq_length = 2048 # Can increase for longer reasoning traces
-    lora_rank = 64 # Larger rank = smarter, but slower
-
     with open('/root/sia/qwq_tokenizer_config.json', 'r') as f:
         tokenizer_config = json.load(f)
 
@@ -84,32 +81,23 @@ def main():
 
     model, _returned_tokenizer = FastLanguageModel.from_pretrained(
         model_name = args.base_model,
-        max_seq_length = max_seq_length,
-        load_in_4bit = True, # False for LoRA 16bit
-        fast_inference = True, # Enable vLLM fast inference
-        max_lora_rank = lora_rank,
-        gpu_memory_utilization = 0.5, # Reduce if out of memory
-        tokenizer = tokenizer,
+        gpu_memory_utilization = 0.5,
     )
 
     model = FastLanguageModel.get_peft_model(
         model,
-        r = lora_rank, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
         target_modules = [
             "q_proj", "k_proj", "v_proj", "o_proj",
             "gate_proj", "up_proj", "down_proj",
         ], # Remove QKVO if out of memory
-        lora_alpha = lora_rank,
         lora_dropout = 0, # Supports any, but = 0 is optimized
         bias = "none",    # Supports any, but = "none" is optimized
         use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
         random_state = 3407,
-        use_rslora = False,  # We support rank stabilized LoRA
-        loftq_config = None, # And LoftQ
     )
 
     training_args = TrainingArguments(
-        output_dir=str(args.output_dir),
+        output_dir=str(args.output_dir) + "_train",
         num_train_epochs=3,
         per_device_train_batch_size=1,
         gradient_accumulation_steps=16,
@@ -133,7 +121,6 @@ def main():
         args=training_args,
         train_dataset=dataset.to_transformers_dataset(tokenizer),
         dataset_text_field="messages",
-        max_seq_length=max_seq_length,
     )
 
     trainer.train()
@@ -141,6 +128,7 @@ def main():
     model.save_pretrained_merged(
         str(args.output_dir), 
         tokenizer=tokenizer,
+        save_method="merged_4bit_forced"
     )
 
 if __name__ == "__main__":