Upload train_qwen3_distillation.py with huggingface_hub

train_qwen3_distillation.py

@@ -0,0 +1,336 @@
+# /// script
+# dependencies = ["transformers>=4.40.0", "datasets", "torch", "accelerate", "peft>=0.7.0", "trackio", "bitsandbytes"]
+# ///
+
+import torch
+import torch.nn.functional as F
+from datasets import load_dataset
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    Trainer,
+    TrainingArguments,
+    DataCollatorForSeq2Seq,
+)
+from peft import LoraConfig, get_peft_model
+import trackio
+from typing import Dict, Optional
+import numpy as np
+
+print("="*50)
+print("Knowledge Distillation: Qwen3-4B -> Qwen3-0.6B")
+print("Method: MiniLLM (Reversed KLD + Teacher Sampling)")
+print("Dataset: TeleQnA (Telecommunications)")
+print("="*50)
+
+# Configuration
+TEACHER_MODEL = "Qwen/Qwen3-4B"
+STUDENT_MODEL = "Qwen/Qwen3-0.6B"
+TEMPERATURE = 2.0  # Temperature for softening distributions
+ALPHA = 0.5  # Weight for distillation loss
+
+# Load tokenizer
+print(f"\nLoading tokenizer from {STUDENT_MODEL}...")
+tokenizer = AutoTokenizer.from_pretrained(STUDENT_MODEL, trust_remote_code=True)
+tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = "right"
+
+# Load TeleQnA dataset
+print("\nLoading TeleQnA dataset...")
+raw_dataset = load_dataset('netop/TeleQnA', split='test')
+
+def format_for_distillation(example):
+    """Convert TeleQnA to chat format"""
+    choices_text = []
+    if 'choices' in example and example['choices']:
+        for i, choice in enumerate(example['choices'], 1):
+            choices_text.append(f'{i}. {choice}')
+
+    question = f"""{example['question']}
+
+Options:
+{chr(10).join(choices_text)}"""
+
+    explanation = example.get('explaination', '') or example.get('explanation', '')
+    answer = f"""{example['answer']}
+
+Explanation: {explanation}"""
+
+    # Create prompt and completion
+    prompt = f"<|im_start|>user\n{question}<|im_end|>\n<|im_start|>assistant\n"
+    completion = f"{answer}<|im_end|>"
+
+    return {"prompt": prompt, "completion": completion}
+
+print("Preprocessing dataset...")
+dataset = raw_dataset.map(format_for_distillation, remove_columns=raw_dataset.column_names)
+
+# Tokenize with prompt/completion structure
+def tokenize_function(examples):
+    # Tokenize prompts (input)
+    prompt_encodings = tokenizer(
+        examples["prompt"],
+        truncation=True,
+        max_length=512,
+        padding=False,
+    )
+
+    # Tokenize completions (target)
+    completion_encodings = tokenizer(
+        examples["completion"],
+        truncation=True,
+        max_length=512,
+        padding=False,
+    )
+
+    # Combine
+    input_ids = [
+        p + c for p, c in zip(prompt_encodings["input_ids"], completion_encodings["input_ids"])
+    ]
+    attention_mask = [
+        p + c for p, c in zip(prompt_encodings["attention_mask"], completion_encodings["attention_mask"])
+    ]
+
+    # Labels: -100 for prompt (don't compute loss), actual tokens for completion
+    labels = [
+        [-100] * len(p) + c for p, c in zip(prompt_encodings["input_ids"], completion_encodings["input_ids"])
+    ]
+
+    return {
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+        "labels": labels,
+    }
+
+print("Tokenizing dataset...")
+tokenized_dataset = dataset.map(
+    tokenize_function,
+    batched=True,
+    remove_columns=["prompt", "completion"],
+)
+
+# Create train/eval split
+print("Creating train/eval split...")
+dataset_split = tokenized_dataset.train_test_split(test_size=0.1, seed=42)
+train_dataset = dataset_split["train"]
+eval_dataset = dataset_split["test"]
+
+print(f"Train examples: {len(train_dataset)}")
+print(f"Eval examples: {len(eval_dataset)}")
+
+# Load Teacher Model (frozen)
+print(f"\nLoading teacher model: {TEACHER_MODEL}...")
+teacher_model = AutoModelForCausalLM.from_pretrained(
+    TEACHER_MODEL,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    trust_remote_code=True,
+)
+teacher_model.eval()
+for param in teacher_model.parameters():
+    param.requires_grad = False
+print("✓ Teacher model loaded and frozen")
+
+# Load Student Model (trainable with LoRA)
+print(f"\nLoading student model: {STUDENT_MODEL}...")
+student_model = AutoModelForCausalLM.from_pretrained(
+    STUDENT_MODEL,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    trust_remote_code=True,
+)
+
+# Apply LoRA
+lora_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"
+)
+student_model = get_peft_model(student_model, lora_config)
+student_model.print_trainable_parameters()
+print("✓ Student model loaded with LoRA")
+
+# MiniLLM Distillation Trainer
+class MiniLLMTrainer(Trainer):
+    """
+    MiniLLM approach with:
+    1. Reversed KL Divergence: KL(student || teacher)
+    2. Teacher token sampling for training targets
+    """
+    def __init__(self, *args, teacher_model=None, temperature=2.0, alpha=0.5, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.teacher_model = teacher_model
+        self.temperature = temperature
+        self.alpha = alpha
+        self.use_teacher_sampling = True  # MiniLLM uses teacher sampling
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        """
+        MiniLLM Loss:
+        1. Sample tokens from teacher distribution
+        2. Compute Reversed KLD between student and teacher
+        3. Combine with cross-entropy loss
+        """
+        input_ids = inputs["input_ids"]
+        attention_mask = inputs["attention_mask"]
+        labels = inputs.pop("labels")
+
+        # Get teacher predictions (no gradient)
+        with torch.no_grad():
+            teacher_outputs = self.teacher_model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+            )
+            teacher_logits = teacher_outputs.logits
+
+            # Teacher token sampling (key part of MiniLLM)
+            if self.use_teacher_sampling and self.training:
+                # Sample from teacher's softmax distribution
+                teacher_probs = F.softmax(teacher_logits / self.temperature, dim=-1)
+                # Sample tokens: [batch, seq_len]
+                sampled_tokens = torch.multinomial(
+                    teacher_probs.view(-1, teacher_probs.size(-1)),
+                    num_samples=1
+                ).view(teacher_probs.size(0), teacher_probs.size(1))
+
+                # Replace labels with teacher-sampled tokens (except where labels are -100)
+                mask = labels != -100
+                labels = torch.where(mask, sampled_tokens, labels)
+
+        # Student forward pass
+        student_outputs = model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+        )
+        student_logits = student_outputs.logits
+
+        # 1. Cross-Entropy Loss (with teacher-sampled tokens)
+        ce_loss = F.cross_entropy(
+            student_logits.view(-1, student_logits.size(-1)),
+            labels.view(-1),
+            ignore_index=-100,
+            reduction='mean'
+        )
+
+        # 2. Reversed KL Divergence: KL(student || teacher)
+        # This encourages student to cover all modes of teacher distribution
+        student_log_probs = F.log_softmax(student_logits / self.temperature, dim=-1)
+        teacher_log_probs = F.log_softmax(teacher_logits / self.temperature, dim=-1)
+        student_probs = F.softmax(student_logits / self.temperature, dim=-1)
+
+        # Reversed KLD = sum(P_student * log(P_student / P_teacher))
+        reversed_kl = torch.sum(
+            student_probs * (student_log_probs - teacher_log_probs),
+            dim=-1
+        )
+
+        # Mask padding and non-target tokens
+        loss_mask = (labels != -100).float()
+        if loss_mask.dim() == 2:
+            # If labels are 2D, add dimension for broadcasting
+            loss_mask = loss_mask.unsqueeze(-1)
+
+        reversed_kl_masked = (reversed_kl * loss_mask.squeeze(-1)).sum() / (loss_mask.sum() + 1e-8)
+
+        # Scale by temperature squared
+        reversed_kl_masked = reversed_kl_masked * (self.temperature ** 2)
+
+        # Combined loss: alpha * Reversed_KL + (1-alpha) * CE
+        total_loss = self.alpha * reversed_kl_masked + (1 - self.alpha) * ce_loss
+
+        # Logging
+        if self.state.global_step % self.args.logging_steps == 0:
+            self.log({
+                "loss/total": total_loss.item(),
+                "loss/reversed_kl": reversed_kl_masked.item(),
+                "loss/cross_entropy": ce_loss.item(),
+            })
+
+        return (total_loss, student_outputs) if return_outputs else total_loss
+
+# Training arguments
+training_args = TrainingArguments(
+    output_dir="qwen3-0.6b-telecom-distilled",
+
+    # Training
+    num_train_epochs=3,
+    per_device_train_batch_size=1,
+    per_device_eval_batch_size=1,
+    gradient_accumulation_steps=16,
+
+    # Optimization
+    learning_rate=2e-4,
+    lr_scheduler_type="cosine",
+    warmup_ratio=0.1,
+    weight_decay=0.01,
+
+    # Evaluation
+    eval_strategy="steps",
+    eval_steps=100,
+    save_strategy="steps",
+    save_steps=200,
+    save_total_limit=3,
+
+    # Logging
+    logging_steps=10,
+    report_to="trackio",
+    run_name="qwen3-0.6b-telecom-minillm",
+
+    # Memory
+    gradient_checkpointing=True,
+    bf16=True,
+
+    # Hub
+    push_to_hub=True,
+    hub_model_id="wlabchoi/qwen3-0.6b-telecom-distilled",
+    hub_strategy="every_save",
+    hub_private_repo=False,
+
+    # Performance
+    dataloader_num_workers=4,
+    remove_unused_columns=False,
+)
+
+# Data collator
+data_collator = DataCollatorForSeq2Seq(
+    tokenizer=tokenizer,
+    model=student_model,
+    padding=True,
+)
+
+# Initialize trainer
+print("\nInitializing MiniLLM Trainer...")
+trainer = MiniLLMTrainer(
+    model=student_model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    data_collator=data_collator,
+    teacher_model=teacher_model,
+    temperature=TEMPERATURE,
+    alpha=ALPHA,
+)
+
+# Start training
+print("\n" + "="*50)
+print("Starting MiniLLM Knowledge Distillation...")
+print(f"✓ Teacher: {TEACHER_MODEL} (frozen)")
+print(f"✓ Student: {STUDENT_MODEL} (LoRA)")
+print(f"✓ Method: Reversed KLD + Teacher Sampling")
+print(f"✓ Temperature: {TEMPERATURE}")
+print(f"✓ Alpha: {ALPHA}")
+print(f"✓ Dataset: TeleQnA ({len(train_dataset)} train, {len(eval_dataset)} eval)")
+print("="*50 + "\n")
+
+trainer.train()
+
+# Push final model
+print("\nPushing distilled model to Hub...")
+trainer.push_to_hub(commit_message="MiniLLM distillation: Qwen3-4B -> Qwen3-0.6B on TeleQnA")
+
+print("\n" + "="*50)
+print("✓ Knowledge Distillation Complete!")
+print("="*50)