Upload train_qwen3_distillation.py with huggingface_hub

train_qwen3_distillation.py

@@ -2,6 +2,7 @@
 # dependencies = ["transformers>=4.40.0", "datasets", "torch", "accelerate", "peft>=0.7.0", "trackio", "bitsandbytes"]
 # ///
 
+import os
 import torch
 import torch.nn.functional as F
 from datasets import load_dataset
@@ -17,6 +18,9 @@ import trackio
 from typing import Dict, Optional
 import numpy as np
 
+# Disable tokenizer parallelism warning
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
 print("="*50)
 print("Knowledge Distillation: Qwen3-4B -> Qwen3-0.6B")
 print("Method: MiniLLM (Reversed KLD + Teacher Sampling)")
@@ -167,7 +171,7 @@ class MiniLLMTrainer(Trainer):
         self.alpha = alpha
         self.use_teacher_sampling = True  # MiniLLM uses teacher sampling
 
-    def compute_loss(self, model, inputs, return_outputs=False):
+    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
         """
         MiniLLM Loss:
         1. Sample tokens from teacher distribution
@@ -290,7 +294,7 @@ training_args = TrainingArguments(
     hub_private_repo=False,
 
     # Performance
-    dataloader_num_workers=4,
+    dataloader_num_workers=0,  # Avoid multiprocessing issues with tokenizers
     remove_unused_columns=False,
 )