Upload train_qwen3_wirelessmath.py with huggingface_hub

train_qwen3_wirelessmath.py

@@ -0,0 +1,140 @@
+# /// script
+# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "datasets", "transformers", "accelerate", "bitsandbytes"]
+# ///
+
+import os
+import torch
+from datasets import load_dataset
+from peft import LoraConfig
+from trl import SFTTrainer, SFTConfig
+import trackio
+
+# Disable tokenizer parallelism warning
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+print("="*50)
+print("Fine-tuning Qwen3-0.6B on WirelessMATHBench-XL")
+print("Method: SFT with LoRA")
+print("Dataset: Wireless Communications Math")
+print("="*50)
+
+# Load WirelessMATHBench-XL dataset
+print("\nLoading WirelessMATHBench-XL dataset...")
+train_dataset = load_dataset('XINLI1997/WirelessMATHBench-XL', split='train')
+eval_dataset = load_dataset('XINLI1997/WirelessMATHBench-XL', split='test')
+
+print(f"Train examples: {len(train_dataset)}")
+print(f"Eval examples: {len(eval_dataset)}")
+
+def format_for_sft(example):
+    """
+    Convert WirelessMATHBench-XL format to chat messages
+    Dataset has: prompt (pre-formatted), correct_answer, and other fields
+    """
+    # Use the pre-formatted prompt
+    prompt = example['prompt']
+    answer = example['correct_answer']
+
+    # Create chat format
+    messages = [
+        {'role': 'user', 'content': prompt},
+        {'role': 'assistant', 'content': answer}
+    ]
+
+    return {'messages': messages}
+
+print("Preprocessing dataset...")
+train_dataset = train_dataset.map(
+    format_for_sft,
+    remove_columns=train_dataset.column_names
+)
+eval_dataset = eval_dataset.map(
+    format_for_sft,
+    remove_columns=eval_dataset.column_names
+)
+
+# Configure LoRA for efficient fine-tuning
+print("\nConfiguring LoRA...")
+peft_config = LoraConfig(
+    r=16,
+    lora_alpha=32,
+    lora_dropout=0.05,
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+    bias="none",
+    task_type="CAUSAL_LM"
+)
+
+# Configure SFT training
+print("Configuring training arguments...")
+training_args = SFTConfig(
+    output_dir="qwen3-wireless-math",
+
+    # Training hyperparameters
+    num_train_epochs=3,
+    per_device_train_batch_size=4,
+    per_device_eval_batch_size=4,
+    gradient_accumulation_steps=4,  # Effective batch size = 16
+
+    # Optimization
+    learning_rate=2e-4,
+    lr_scheduler_type="cosine",
+    warmup_ratio=0.1,
+    weight_decay=0.01,
+
+    # Evaluation and saving
+    eval_strategy="steps",
+    eval_steps=100,
+    save_strategy="steps",
+    save_steps=200,
+    save_total_limit=3,
+
+    # Logging and monitoring
+    logging_steps=10,
+    report_to="trackio",
+    run_name="qwen3-0.6b-wireless-math",
+    project="wireless-math-finetuning",
+
+    # Memory optimization
+    gradient_checkpointing=False,  # Disabled to avoid gradient computation issues
+    bf16=True,
+
+    # Hub integration
+    push_to_hub=True,
+    hub_model_id="wlabchoi/qwen3-0.6b-wireless-math",
+    hub_strategy="every_save",
+    hub_private_repo=False,
+
+    # Performance
+    dataloader_num_workers=0,  # Avoid multiprocessing issues
+    remove_unused_columns=False,
+)
+
+# Initialize trainer
+print("\nInitializing SFT Trainer...")
+trainer = SFTTrainer(
+    model="Qwen/Qwen3-0.6B",
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    peft_config=peft_config,
+    args=training_args,
+)
+
+# Start training
+print("\n" + "="*50)
+print("Starting Fine-Tuning...")
+print(f"Model: Qwen3-0.6B")
+print(f"Dataset: WirelessMATHBench-XL")
+print(f"Train: {len(train_dataset)} examples")
+print(f"Eval: {len(eval_dataset)} examples")
+print(f"Epochs: 3")
+print("="*50 + "\n")
+
+trainer.train()
+
+# Push final model to Hub
+print("\nPushing final model to Hub...")
+trainer.push_to_hub(commit_message="Fine-tuning complete - Qwen3-0.6B on WirelessMATHBench-XL")
+
+print("\n" + "="*50)
+print("Fine-Tuning Completed Successfully!")
+print("="*50)