import json import os import copy import random from typing import Any, Dict, List, Optional, Tuple, Union import torch import torch.nn.functional as F import yaml from torch.utils.data import Dataset from transformers import AutoTokenizer, AutoModel, Trainer, TrainingArguments from peft import LoraConfig, get_peft_model def load_config(config_path: str) -> Dict: with open(config_path, "r", encoding="utf-8") as f: config = yaml.safe_load(f) return config def get_deepspeed_config(config: Dict[str, Any]) -> Dict[str, Any]: """Creating a DeepSpeed Configuration""" return { "train_batch_size": "auto", "train_micro_batch_size_per_gpu": "auto", "gradient_accumulation_steps": "auto", "gradient_clipping": "auto", "zero_allow_untested_optimizer": True, "bf16": { "enabled": "auto" }, "zero_optimization": { "stage": 2, "allgather_partitions": True, "allgather_bucket_size": 2e8, "reduce_scatter": True, "reduce_bucket_size": 2e8, "overlap_comm": True, "contiguous_gradients": True, }, } def prepare_models(config: Dict[str, Any]): """Prepare Student and Frozen Teacher models""" torch_dtype = config['model']['torch_dtype'] model_name = config['model']['name'] trust_remote_code = config['model']['trust_remote_code'] tokenizer = AutoTokenizer.from_pretrained( model_name, trust_remote_code=trust_remote_code, ) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token # 1. Load base model as Student student_model = AutoModel.from_pretrained( model_name, torch_dtype=getattr(torch, torch_dtype) if isinstance(torch_dtype, str) else torch_dtype, trust_remote_code=trust_remote_code, ) # 2. Deep copy base model as Teacher and freeze parameters teacher_model = copy.deepcopy(student_model) for param in teacher_model.parameters(): param.requires_grad = False teacher_model.eval() # 3. Inject LoRA into Student model lora_config = LoraConfig( r=config['lora']['r'], lora_alpha=config['lora']['lora_alpha'], target_modules=config['lora']['target_modules'], lora_dropout=config['lora']['lora_dropout'], bias=config['lora']['bias'], task_type=config['lora']['task_type'], ) student_model = get_peft_model(student_model, lora_config) student_model.print_trainable_parameters() return student_model, teacher_model, tokenizer class TrajectoryDataset(Dataset): def __init__(self, data_path: str, tokenizer: AutoTokenizer, delta: int = 4, mask_token_id: int = 126336, sample_ratio: float = 1.0, shuffle: bool = True, seed: int = 42): self.data_path = data_path self.tokenizer = tokenizer self.mask_token_id = mask_token_id self.delta = delta self.sample_ratio = max(0.0, min(1.0, float(sample_ratio))) self.shuffle = shuffle self.seed = seed self.data = [] with open(self.data_path, "r", encoding="utf-8") as f: for line in f: line = line.strip() if not line: continue self.data.append(json.loads(line)) self.step_keys = [None] * len(self.data) self.index_map = [] for record_idx, record in enumerate(self.data): trajectory = record.get("trajectory", {}) # step0 (all Mask) -> stepN (Final answer) step_keys = sorted(trajectory.keys(), key=self._step_key_to_idx) if len(step_keys) < 2: continue self.step_keys[record_idx] = step_keys # Exclude steps that cannot perform delta prediction for step_idx in range(len(step_keys) - 1): self.index_map.append((record_idx, step_idx)) if self.sample_ratio < 1.0: target_size = int(len(self.index_map) * self.sample_ratio) if target_size < len(self.index_map): self.index_map = random.sample(self.index_map, target_size) # Shuffle index_map globally to randomize training data order if self.shuffle: rng = random.Random(self.seed) rng.shuffle(self.index_map) def __len__(self): return len(self.index_map) @staticmethod def _step_key_to_idx(key: str) -> int: digits = "".join([c for c in key if c.isdigit()]) return int(digits) if digits else 0 def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: record_idx, step_idx = self.index_map[idx] record = self.data[record_idx] prompt = record["prompt"] trajectory = record["trajectory"] answer_gt = record.get("groundtruth", "") step_keys = self.step_keys[record_idx] # Determine skip-step target (later steps have fewer masks) target_idx = min(step_idx + self.delta, len(step_keys) - 1) x_t = trajectory[step_keys[step_idx]] x_target = trajectory[step_keys[target_idx]] # 1. Build Student (blind guess) input student_prompt_ids = self.tokenizer(prompt, add_special_tokens=False).input_ids student_input_ids = student_prompt_ids + x_t # 2. Teacher Inputs if prompt.startswith("<|startoftext|>"): system_prompt = f"<|startoftext|><|start_header_id|>system<|end_header_id|>\n\nReference Answer: {answer_gt}<|eot_id|>" prompt_without_bos = prompt[len("<|startoftext|>"):] teacher_prompt_text = system_prompt + prompt_without_bos else: system_prompt = f"<|start_header_id|>system<|end_header_id|>\n\nReference Answer: {answer_gt}<|eot_id|>" teacher_prompt_text = system_prompt + prompt teacher_prompt_ids = self.tokenizer(teacher_prompt_text, add_special_tokens=False).input_ids teacher_input_ids = teacher_prompt_ids + x_t # 3. Build hard labels (CE Target) # Only provide labels for the answer part; prompt part uses -100 to ignore labels_tail = [] ce_mask_tail = [] # CE region: currently MASK and Target is not MASK (has hard label) kl_mask_tail = [] # KL region: currently MASK and Target is still MASK (no hard label, needs Teacher guidance) for tk_t, tk_target in zip(x_t, x_target): if tk_t == self.mask_token_id and tk_target != self.mask_token_id: # CE region: currently MASK, Target revealed -> train with hard label labels_tail.append(tk_target) ce_mask_tail.append(1.0) kl_mask_tail.append(0.0) elif tk_t == self.mask_token_id and tk_target == self.mask_token_id: # KL region: currently MASK, Target still MASK -> use Teacher soft label labels_tail.append(-100) ce_mask_tail.append(0.0) kl_mask_tail.append(1.0) else: # Non-MASK position: not involved in any loss computation labels_tail.append(-100) ce_mask_tail.append(0.0) kl_mask_tail.append(0.0) # Concatenate prompt part (prompt part does not contribute to loss) labels = [-100] * len(student_prompt_ids) + labels_tail student_ce_mask = [0.0] * len(student_prompt_ids) + ce_mask_tail student_kl_mask = [0.0] * len(student_prompt_ids) + kl_mask_tail teacher_kl_mask = [0.0] * len(teacher_prompt_ids) + kl_mask_tail return { "student_input_ids": torch.tensor(student_input_ids, dtype=torch.long), "teacher_input_ids": torch.tensor(teacher_input_ids, dtype=torch.long), "labels": torch.tensor(labels, dtype=torch.long), "student_ce_mask": torch.tensor(student_ce_mask, dtype=torch.float), "student_kl_mask": torch.tensor(student_kl_mask, dtype=torch.float), "teacher_kl_mask": torch.tensor(teacher_kl_mask, dtype=torch.float), } def build_collator(tokenizer: AutoTokenizer): pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0 def collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]: student_input_ids = [item["student_input_ids"] for item in batch] teacher_input_ids = [item["teacher_input_ids"] for item in batch] labels = [item["labels"] for item in batch] student_ce_mask = [item["student_ce_mask"] for item in batch] student_kl_mask = [item["student_kl_mask"] for item in batch] teacher_kl_mask = [item["teacher_kl_mask"] for item in batch] from torch.nn.utils.rnn import pad_sequence # Right Padding student_batch = pad_sequence(student_input_ids, batch_first=True, padding_value=pad_id) teacher_batch = pad_sequence(teacher_input_ids, batch_first=True, padding_value=pad_id) labels_batch = pad_sequence(labels, batch_first=True, padding_value=-100) student_ce_mask_batch = pad_sequence(student_ce_mask, batch_first=True, padding_value=0.0) student_kl_mask_batch = pad_sequence(student_kl_mask, batch_first=True, padding_value=0.0) teacher_kl_mask_batch = pad_sequence(teacher_kl_mask, batch_first=True, padding_value=0.0) student_att_mask = (student_batch != pad_id).long() teacher_att_mask = (teacher_batch != pad_id).long() return { "student_input_ids": student_batch, "student_attention_mask": student_att_mask, "teacher_input_ids": teacher_batch, "teacher_attention_mask": teacher_att_mask, "labels": labels_batch, "student_ce_mask": student_ce_mask_batch, "student_kl_mask": student_kl_mask_batch, "teacher_kl_mask": teacher_kl_mask_batch, } return collate_fn class DLMTrainer(Trainer): def __init__(self, teacher_model, lambda_ce=1.0, lambda_kl=1.0, tau=1.0, mask_token_id=126336, **kwargs): super().__init__(**kwargs) self.teacher_model = teacher_model self.lambda_ce = lambda_ce self.lambda_kl = lambda_kl self.tau = tau self.mask_token_id = mask_token_id def compute_loss(self, model, inputs, return_outputs=False, **kwargs): # Ensure Teacher model is on the correct GPU (DeepSpeed ZeRO-2 compatible) if next(self.teacher_model.parameters()).device != model.device: self.teacher_model = self.teacher_model.to(model.device) # ================== 1. Student Forward ================== student_outputs = model( input_ids=inputs["student_input_ids"], attention_mask=inputs["student_attention_mask"] ) student_logits = student_outputs.logits # ================== 2. Teacher Forward ================== with torch.no_grad(): teacher_outputs = self.teacher_model( input_ids=inputs["teacher_input_ids"], attention_mask=inputs["teacher_attention_mask"] ) teacher_logits = teacher_outputs.logits # ================== 3. Extract CE and KL region Logits ================== # CE region: currently MASK and Target revealed (has hard label) s_ce_bool = inputs["student_ce_mask"].bool() ce_student_logits = student_logits[s_ce_bool] ce_labels = inputs["labels"][s_ce_bool] # KL region: currently MASK and Target still MASK (no hard label, needs Teacher guidance) s_kl_bool = inputs["student_kl_mask"].bool() t_kl_bool = inputs["teacher_kl_mask"].bool() kl_student_logits = student_logits[s_kl_bool] kl_teacher_logits = teacher_logits[t_kl_bool] assert kl_student_logits.shape[0] == kl_teacher_logits.shape[0], \ f"KL mask mismatch: student {kl_student_logits.shape[0]} vs teacher {kl_teacher_logits.shape[0]}" # ================== 4. Dual Loss (spatially exclusive) ================== # a. Trajectory CE Loss (only on mask tokens with hard labels) if ce_student_logits.numel() > 0 and (ce_labels != -100).any(): loss_ce = F.cross_entropy(ce_student_logits, ce_labels, ignore_index=-100) else: loss_ce = torch.tensor(0.0, device=model.device) # b. Privileged KL Loss (only on mask tokens without hard labels, still masked) if kl_student_logits.numel() > 0: p_teacher = F.softmax(kl_teacher_logits / self.tau, dim=-1) log_p_student = F.log_softmax(kl_student_logits / self.tau, dim=-1) loss_kl = F.kl_div(log_p_student, p_teacher, reduction='batchmean') * (self.tau ** 2) else: loss_kl = torch.tensor(0.0, device=model.device) # total Loss if ce_student_logits.numel() == 0 and kl_student_logits.numel() == 0: loss = torch.tensor(0.0, device=model.device, requires_grad=True) else: loss = self.lambda_ce * loss_ce + self.lambda_kl * loss_kl self.log({ "loss_ce": float(loss_ce.detach()), "loss_kl": float(loss_kl.detach()), "ce_tokens": int(s_ce_bool.sum()), "kl_tokens": int(s_kl_bool.sum()), }) return (loss, student_outputs) if return_outputs else loss def main(): config = load_config("configs/config_llada.yaml") training_args = TrainingArguments( **config['training'], deepspeed=get_deepspeed_config(config), ddp_find_unused_parameters=False, remove_unused_columns=False, ) # Save config to output_dir for reference output_dir = config['training']['output_dir'] os.makedirs(output_dir, exist_ok=True) config_save_path = os.path.join(output_dir, "config_used.yaml") with open(config_save_path, "w", encoding="utf-8") as f: yaml.dump(config, f, default_flow_style=False, allow_unicode=True) print(f"Config saved to: {config_save_path}") # Prepare student model and teacher model student_model, teacher_model, tokenizer = prepare_models(config) sample_ratio = config.get("data", {}).get("sample_ratio", 1.0) shuffle = config.get("data", {}).get("shuffle", True) seed = config.get("data", {}).get("seed", 42) train_dataset = TrajectoryDataset( data_path="/home/Traj-Distill/data/LLaDA_data.jsonl", tokenizer=tokenizer, delta=8, mask_token_id=126336, sample_ratio=sample_ratio, shuffle=shuffle, seed=seed, ) trainer = DLMTrainer( model=student_model, teacher_model=teacher_model, args=training_args, train_dataset=train_dataset, data_collator=build_collator(tokenizer), mask_token_id=126336, lambda_ce=1.0, lambda_kl=1.0, tau=1.0, ) trainer.train() if __name__ == "__main__": main()