diff --git a/cookbook/rl/mopd/mopd_ctkd.py b/cookbook/rl/mopd/mopd_ctkd.py new file mode 100644 index 000000000..c0d1ac163 --- /dev/null +++ b/cookbook/rl/mopd/mopd_ctkd.py @@ -0,0 +1,415 @@ +import os +from typing import List, Optional + +import torch +from peft import LoraConfig + +import twinkle +from twinkle import DeviceMesh, DeviceGroup, get_device_placement, get_logger +from twinkle.data_format import SamplingParams +from twinkle.dataloader import DataLoader +from twinkle.dataset import Dataset, DatasetMeta +from twinkle.loss import CTKDLoss +from twinkle.model import TransformersModel +from twinkle.sampler import vLLMSampler + +logger = get_logger() + +# ── Configuration ───────────────────────────────────────────────────────────── +STUDENT_MODEL_ID = os.environ.get('STUDENT_MODEL_ID', '/nas/disk1/qwen2.5-0.5b-instruct') +TEACHER_MODEL_ID = os.environ.get('TEACHER_MODEL_ID', '/model/Qwen3-0.6B') +DATASET_ID = os.environ.get('DATASET_ID', '/model/liujihui/twinkle_client_st/httpserver/models/DG04F8511A00100002/messages.jsonl') + +MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 1)) +SAMPLER_GPUS = int(os.environ.get('SAMPLER_GPUS', 1)) +SHARED_TEACHER_GPUS = bool(os.environ.get('SHARED_TEACHER_GPUS', False)) +NUM_GPUS = 2 + +BATCH_SIZE = int(os.environ.get('BATCH_SIZE', 8)) +MAX_STEPS = int(os.environ.get('MAX_STEPS', 10)) +LEARNING_RATE = float(os.environ.get('LR', 1e-5)) +GRADIENT_ACCUMULATION_STEPS = int(os.environ.get('GRADIENT_ACCUMULATION_STEPS', 4)) + +CTKD_TEMPERATURE = float(os.environ.get('CTKD_TEMPERATURE', 0.8)) +CTKD_MAX_LENGTH = int(os.environ.get('CTKD_MAX_LENGTH', 4)) +CTKD_BETA = float(os.environ.get('CTKD_BETA', 0.95)) +CTKD_GAMMA = float(os.environ.get('CTKD_GAMMA', 0.1)) +CTKD_LOSS_TYPE = os.environ.get('CTKD_LOSS_TYPE', 'pkl') +CTKD_TOPK = int(os.environ.get('CTKD_TOPK', 512)) + +ADAPTER_NAME = 'default' +MAX_LENGTH = int(os.environ.get('MAX_LENGTH', 2048)) +MAX_NEW_TOKENS = int(os.environ.get('MAX_NEW_TOKENS', 2048)) +N_SAMPLES = int(os.environ.get('N_SAMPLES', 1)) +SHARED_TEACHER_GPUS = bool(os.environ.get('SHARED_TEACHER_GPUS', False)) + + +# ── Utility ─────────────────────────────────────────────────────────────────── + +def convert_topk_prompt_logprobs( + topk_prompt_logprobs_batch: List[List[Optional[List[tuple]]]], + topk: int = 64, +) -> dict: + batch_logprobs = [] + batch_indices = [] + + for seq_topk in topk_prompt_logprobs_batch: + seq_logprobs = [] + seq_indices = [] + for pos_topk in seq_topk: + if pos_topk is None: + seq_logprobs.append([0.0] * topk) + seq_indices.append([0] * topk) + else: + seq_logprobs.append([lp for _, lp in pos_topk]) + seq_indices.append([tid for tid, _ in pos_topk]) + batch_logprobs.append(seq_logprobs) + batch_indices.append(seq_indices) + + max_len = max(len(seq) for seq in batch_logprobs) if batch_logprobs else 1 + for i in range(len(batch_logprobs)): + pad_len = max_len - len(batch_logprobs[i]) + if pad_len > 0: + batch_logprobs[i].extend([[0.0] * topk] * pad_len) + batch_indices[i].extend([[0] * topk] * pad_len) + + # Roll to align with labels (first position has no valid logprobs) + return { + 'teacher_topk_logprobs': torch.roll(torch.tensor(batch_logprobs, dtype=torch.float32), shifts=-1, dims=1), + 'teacher_topk_indices': torch.roll(torch.tensor(batch_indices, dtype=torch.long), shifts=-1, dims=1), + } + + +def align_teacher_logprobs_to_student( + teacher_topk_logprobs: torch.Tensor, + teacher_topk_indices: torch.Tensor, + teacher_input_ids: torch.Tensor, + student_input_ids_list: List[torch.Tensor], + student_tokenizer, + teacher_tokenizer, + topk: int = 512, + pad_token_id: int = 0, +) -> dict: + batch_size = len(student_input_ids_list) + max_student_len = max(len(ids) for ids in student_input_ids_list) + + aligned_logprobs = torch.zeros(batch_size, max_student_len, topk, dtype=teacher_topk_logprobs.dtype) + aligned_indices = torch.zeros(batch_size, max_student_len, topk, dtype=teacher_topk_indices.dtype) + + for b in range(batch_size): + student_ids = student_input_ids_list[b] + student_len = len(student_ids) + + # Build student character spans by incremental decoding + student_char_spans = [] + for pos in range(student_len): + prefix_text = student_tokenizer.decode(student_ids[:pos+1].tolist(), skip_special_tokens=False) + if pos == 0: + start_char = 0 + else: + prev_prefix = student_tokenizer.decode(student_ids[:pos].tolist(), skip_special_tokens=False) + start_char = len(prev_prefix) + end_char = len(prefix_text) + student_char_spans.append((start_char, end_char)) + + # Find actual teacher sequence length (before padding) + teacher_seq_len = teacher_input_ids.shape[1] + for t in range(teacher_input_ids.shape[1] - 1, -1, -1): + if teacher_input_ids[b, t].item() != pad_token_id: + teacher_seq_len = t + 1 + break + + # Build teacher character spans + teacher_char_spans = [] + teacher_ids_list = teacher_input_ids[b, :teacher_seq_len].tolist() + for pos in range(teacher_seq_len): + prefix_text = teacher_tokenizer.decode(teacher_ids_list[:pos+1], skip_special_tokens=False) + if pos == 0: + start_char = 0 + else: + prev_prefix = teacher_tokenizer.decode(teacher_ids_list[:pos], skip_special_tokens=False) + start_char = len(prev_prefix) + end_char = len(prefix_text) + teacher_char_spans.append((start_char, end_char)) + + # For each student position, find the best-matching teacher position + for s_pos in range(student_len): + s_start, s_end = student_char_spans[s_pos] + + best_t_pos = -1 + best_overlap = 0 + for t_pos in range(teacher_seq_len): + t_start, t_end = teacher_char_spans[t_pos] + overlap = min(s_end, t_end) - max(s_start, t_start) + if overlap > best_overlap: + best_overlap = overlap + best_t_pos = t_pos + + if best_t_pos >= 0 and best_t_pos < teacher_topk_logprobs.shape[1]: + aligned_logprobs[b, s_pos] = teacher_topk_logprobs[b, best_t_pos] + aligned_indices[b, s_pos] = teacher_topk_indices[b, best_t_pos] + + return { + 'teacher_topk_logprobs': aligned_logprobs, + 'teacher_topk_indices': aligned_indices, + } +def create_dataset(): + """创建用于蒸馏的全文(prompt + response)数据集。 + + 数据集使用 student tokenizer 编码。Teacher 会将文本解码后， + 使用自己的 tokenizer 重新编码，以实现跨 tokenizer 知识蒸馏。 + """ + dataset = Dataset(DatasetMeta(DATASET_ID, data_slice=range(10000))) + dataset.set_template('Template', model_id=STUDENT_MODEL_ID, max_length=MAX_LENGTH) + dataset.encode(load_from_cache_file=True) + return dataset + + +# ── Training ────────────────────────────────────────────────────────────────── + +def train(): + import time + start_time = time.perf_counter() + print('记录开始时间') + + # Initialize device groups based on shared mode + if SHARED_TEACHER_GPUS: + device_groups = [ + DeviceGroup(name='student_model', ranks=MODEL_GPUS, device_type='npu'), + DeviceGroup(name='teacher_sampler', ranks=SAMPLER_GPUS, device_type='npu'), + ] + else: + device_groups = [ + DeviceGroup(name='student_model', ranks=MODEL_GPUS, device_type='npu'), + DeviceGroup(name='teacher_sampler', ranks=SAMPLER_GPUS, device_type='npu'), + ] + + model_mesh = DeviceMesh.from_sizes(world_size=MODEL_GPUS, dp_size=MODEL_GPUS) + sampler_mesh = DeviceMesh.from_sizes(world_size=SAMPLER_GPUS, dp_size=SAMPLER_GPUS) + + twinkle.initialize( + mode='ray', + nproc_per_node=NUM_GPUS, + groups=device_groups, + ) + end_time = time.perf_counter() + elapsed = end_time - start_time + print(f"代码initialize执行耗时: {elapsed:.6f} 秒") + start_time = end_time + + # ── Student model (trainable) ────────────────────────────────────────────── + student_model = TransformersModel( + model_id=STUDENT_MODEL_ID, + device_mesh=model_mesh, + remote_group='student_model', + ) + end_time = time.perf_counter() + elapsed = end_time - start_time + print(f"代码student_model: {elapsed:.6f} 秒") + start_time = end_time + + # LoRA configuration for efficient fine-tuning + lora_config = LoraConfig( + r=8, + lora_alpha=32, + lora_dropout=0.05, + target_modules='all-linear', + ) + student_model.add_adapter_to_model(ADAPTER_NAME, lora_config, gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS) + student_model.set_optimizer('AdamW', lr=LEARNING_RATE, weight_decay=0.01) + student_model.set_lr_scheduler('CosineAnnealingLR', T_max=MAX_STEPS, eta_min=LEARNING_RATE * 0.1) + + # ── Configure CTKDLoss ───────────────────────────────────────────────────── + from transformers import AutoTokenizer + student_tokenizer = AutoTokenizer.from_pretrained(STUDENT_MODEL_ID, trust_remote_code=True) + teacher_tokenizer = AutoTokenizer.from_pretrained(TEACHER_MODEL_ID, trust_remote_code=True) + end_time = time.perf_counter() + elapsed = end_time - start_time + print(f"代码AutoTokenizer: {elapsed:.6f} 秒") + start_time = end_time + + loss_fn = CTKDLoss( + student_tokenizer=student_tokenizer, + teacher_tokenizer_group=[teacher_tokenizer], + max_length=CTKD_MAX_LENGTH, + beta=CTKD_BETA, + gamma=CTKD_GAMMA, + loss_type=CTKD_LOSS_TYPE, + temperature=CTKD_TEMPERATURE, + device=torch.device('npu:0'), + ) + student_model.set_loss(loss_fn, adapter_name=ADAPTER_NAME) + student_model.set_template('QwenTemplate', model_id=STUDENT_MODEL_ID, adapter_name=ADAPTER_NAME) + end_time = time.perf_counter() + elapsed = end_time - start_time + print(f"代码loss_fn: {elapsed:.6f} 秒") + start_time = end_time + + # Log configuration + logger.info(f'GPU Configuration: MODEL_GPUS={MODEL_GPUS}, SAMPLER_GPUS={SAMPLER_GPUS}, SHARED_TEACHER_GPUS={SHARED_TEACHER_GPUS}') + logger.info(f'Total GPUs required: {NUM_GPUS}') + + # Log projection matrix statistics with validation + stats = loss_fn.get_mapping_statistics() + logger.info(f'CTKD Projection Matrix Statistics: {stats}') + + # Validate vocabulary coverage + coverage_ratio = stats['exact_matched'] / stats['total_student_tokens'] + logger.info(f'Vocabulary coverage ratio: {coverage_ratio:.2%}') + if coverage_ratio < 0.3: + logger.warning(f"Low vocabulary coverage ({coverage_ratio:.2%}), consider using models with similar tokenizers") + logger.warning("This may cause poor distillation performance and gradient issues") + + # ── Teacher vLLM samplers ────────────────────────────────────────────────── + if SHARED_TEACHER_GPUS: + teacher_sampler = vLLMSampler( + model_id=TEACHER_MODEL_ID, + engine_args={ + 'gpu_memory_utilization': 0.75, + 'max_model_len': 4096, + 'logprobs_mode': 'raw_logprobs', + 'max_logprobs': CTKD_TOPK, + }, + device_mesh=sampler_mesh, + remote_group='teacher_sampler', + instance_id='teacher_1' + ) + teacher_sampler.set_template('QwenTemplate', model_id=TEACHER_MODEL_ID) + else: + teacher_sampler = vLLMSampler( + model_id=TEACHER_MODEL_ID, + engine_args={ + 'gpu_memory_utilization': 0.75, + 'max_model_len': 4096, + 'logprobs_mode': 'raw_logprobs', + 'max_logprobs': CTKD_TOPK, + }, + device_mesh=sampler_mesh, + remote_group='teacher_sampler', + ) + teacher_sampler.set_template('QwenTemplate', model_id=TEACHER_MODEL_ID) + + # ── DataLoader ───────────────────────────────────────────────────────────── + dataloader = DataLoader( + dataset=create_dataset(), + batch_size=BATCH_SIZE, + min_batch_size=BATCH_SIZE, + device_mesh=model_mesh, + remote_group='student_model', + ) + + # ── Training Loop ────────────────────────────────────────────────────────── + optim_step = 0 + for batch in dataloader: + if optim_step >= MAX_STEPS: + break + if callable(batch): + batch = batch() + + # ── Step 1: Decode student tokens to text for teacher ────────────────── + from twinkle.data_format import Trajectory + + teacher_inputs = [] + student_input_ids_list = [] + for item in batch: + text = student_tokenizer.decode(item['input_ids'], skip_special_tokens=False) + teacher_inputs.append({'messages': [{'role': 'user', 'content': text}]}) + student_input_ids_list.append(torch.tensor(item['input_ids'])) + + # ── Step 2: Teacher computes top-k logprobs ──────────────────────────── + # CRITICAL FIX: Use max_tokens=0 to compute prompt_logprobs only, not generate new tokens + teacher_response = teacher_sampler.sample( + teacher_inputs, + SamplingParams(max_tokens=0, temperature=1.0, prompt_logprobs=CTKD_TOPK), + ) + + # ── Step 3: Convert teacher responses ────────────────────────────────── + teacher_input_data = [seq.new_input_feature for resp in teacher_response for seq in resp.sequences] + + # ── Step 4: Prepare teacher output with alignment ────────────────────── + # 4a. Convert topk to tensor format (teacher's sequence length) + topk_data = convert_topk_prompt_logprobs( + [resp.topk_prompt_logprobs for resp in teacher_response], + topk=CTKD_TOPK, + ) + + # 4b. Get teacher input_ids for alignment + import torch.nn.utils.rnn as rnn_utils + teacher_input_ids_list = [torch.tensor(item['input_ids']) for item in teacher_input_data] + teacher_input_ids = rnn_utils.pad_sequence(teacher_input_ids_list, batch_first=True) + + # 4c. CRITICAL FIX: Align teacher logprobs to student token positions + # This ensures position-wise KL divergence is computed on semantically + # corresponding tokens, not just same index positions. + aligned_topk = align_teacher_logprobs_to_student( + teacher_topk_logprobs=topk_data['teacher_topk_logprobs'], + teacher_topk_indices=topk_data['teacher_topk_indices'], + teacher_input_ids=teacher_input_ids, + student_input_ids_list=student_input_ids_list, + student_tokenizer=student_tokenizer, + teacher_tokenizer=teacher_tokenizer, + topk=CTKD_TOPK, + ) + + # 4d. Create teacher labels aligned to student sequence length + # For CTKD, we want to distill over ALL token positions (not just response). + # So we use input_ids as labels (no -100 masking) instead of student's labels + # which masks prompt positions with -100. + student_labels_list = [] + for item in batch: + # Use input_ids as labels for CTKD - all positions should participate in loss + # This is different from standard supervised training where prompt is masked + labels = torch.tensor(item['input_ids']) + student_labels_list.append(labels) + student_labels = rnn_utils.pad_sequence(student_labels_list, batch_first=True, padding_value=-100) + + # DEBUG: Log teacher_labels statistics + print(f'[CTKD DEBUG] mopd_ctkd: student_labels shape: {student_labels.shape}') + print(f'[CTKD DEBUG] mopd_ctkd: student_labels -100 count: {(student_labels == -100).sum().item()}') + print(f'[CTKD DEBUG] mopd_ctkd: student_labels non -100 count: {(student_labels != -100).sum().item()}') + print(f'[CTKD DEBUG] mopd_ctkd: student_labels first sample first 20 values: {student_labels[0, :20].tolist()}') + + # DEBUG: Log teacher topk data statistics + print(f'[CTKD DEBUG] mopd_ctkd: teacher_topk_logprobs shape: {aligned_topk["teacher_topk_logprobs"].shape}') + print(f'[CTKD DEBUG] mopd_ctkd: teacher_topk_logprobs range: [{aligned_topk["teacher_topk_logprobs"].min().item():.4f}, {aligned_topk["teacher_topk_logprobs"].max().item():.4f}]') + print(f'[CTKD DEBUG] mopd_ctkd: teacher_topk_indices shape: {aligned_topk["teacher_topk_indices"].shape}') + print(f'[CTKD DEBUG] mopd_ctkd: teacher_topk_indices range: [{aligned_topk["teacher_topk_indices"].min().item()}, {aligned_topk["teacher_topk_indices"].max().item()}]') + print(f'[CTKD DEBUG] mopd_ctkd: teacher_input_ids shape: {teacher_input_ids.shape}') + print(f'[CTKD DEBUG] mopd_ctkd: teacher_input_ids first sample first 20 values: {teacher_input_ids[0, :20].tolist()}') + + # Create teacher_output dict with aligned data + teacher_output = { + 'teacher_labels': [student_labels], # All positions participate in CTKD loss + 'teacher_input_ids': [teacher_input_ids], + 'teacher_topk_logprobs_group': [aligned_topk['teacher_topk_logprobs']], + 'teacher_topk_indices_group': [aligned_topk['teacher_topk_indices']], + } + + # ── Step 5: Student forward + CTKD backward ──────────────────────────── + student_model.forward_backward( + inputs=batch, + adapter_name=ADAPTER_NAME, + return_logits=True, + **teacher_output, + ) + + student_model.clip_grad_and_step(adapter_name=ADAPTER_NAME) + + # 5. Logging + if optim_step > 0 and optim_step % 2 == 0: + metric = student_model.calculate_metric(is_training=True, adapter_name=ADAPTER_NAME) + logger.info(f'[Step {optim_step}/{MAX_STEPS}] {metric}') + + # ── Checkpoint ───────────────────────────────────────────────────────── + if optim_step > 0 and optim_step % 100 == 0: + student_model.save(f'mopd-ctkd-ckpt-{optim_step}', adapter_name=ADAPTER_NAME) + + optim_step += 1 + + # Save final checkpoint + student_model.save('mopd-ctkd-final', adapter_name=ADAPTER_NAME) + logger.info('MOPD CTKD training completed.') + + +if __name__ == '__main__': + train() \ No newline at end of file diff --git a/src/twinkle/loss/__init__.py b/src/twinkle/loss/__init__.py index 8e1d0e2ad..59fdf8dfe 100644 --- a/src/twinkle/loss/__init__.py +++ b/src/twinkle/loss/__init__.py @@ -7,6 +7,7 @@ from .grpo import BNPOLoss, CISPOLoss, DRGRPOLoss, GRPOLoss, GSPOLoss, SAPOLoss from .infonce import InfonceLoss from .mse import MSELoss +from .ctkd import CTKDLoss torch_loss_mapping = { 'mse': MSELoss, @@ -28,4 +29,5 @@ 'orpo': ORPOLoss, # Embedding / contrastive losses 'infonce': InfonceLoss, + 'ctkd': CTKDLoss } diff --git a/src/twinkle/loss/ctkd.py b/src/twinkle/loss/ctkd.py new file mode 100644 index 000000000..bfe38efc3 --- /dev/null +++ b/src/twinkle/loss/ctkd.py @@ -0,0 +1,982 @@ +from typing import TYPE_CHECKING, Dict, Optional, Tuple +import hashlib +import pickle + +import torch +import torch.nn.functional as F + +from twinkle.data_format import LossOutput +from twinkle.loss.base import Loss +import threading + +if TYPE_CHECKING: + from transformers import PreTrainedTokenizer + +# Global cache for projection matrices to avoid recomputation +_PROJECTION_MATRIX_CACHE = {} + + +class CTKDLoss(Loss): + """Cross-Tokenizer Knowledge Distillation Loss using X-Token projection matrix. + + This loss enables knowledge distillation between models with different tokenizers + by constructing sparse projection matrices W that align student and teacher + vocabulary spaces for multiple teachers. + + The projection matrix construction follows three steps: + 1. Initialize W[s,t] = 0 for all student token s and teacher token t + 2. Exact match: If student token text equals teacher token text, set W[s,t] = 1 + 3. Multi-token decoding: For unmatched student tokens, decode to text and + re-encode with teacher tokenizer. If the resulting sequence length < L (max_length), + assign weights: W[s, t[i]] = β * γ^i + + Args: + student_tokenizer: Tokenizer for the student model. + teacher_tokenizer_group: List of tokenizers for the teacher models. + teacher_weights: Optional list of weights for each teacher (default: equal weights). + max_length: Maximum span length L for multi-token matching (default: 4). + beta: Base weight β for projection (default: 0.9). + gamma: Decay rate γ for multi-token weights (default: 0.1). + loss_type: Type of KL loss to use - 'pkl' for P-KL or 'hkl' for H-KL (default: 'pkl'). + temperature: Temperature for softmax in KL divergence (default: 1.0). + device: Device to place the projection matrices on (default: None). + """ + + def __init__( + self, + student_tokenizer: 'PreTrainedTokenizer', + teacher_tokenizer_group: list, # List of teacher tokenizers + teacher_weights: Optional[list] = None, # Optional weights for each teacher + max_length: int = 4, + beta: float = 0.9, + gamma: float = 0.1, + loss_type: str = None, # Auto-select based on vocabulary coverage + temperature: float = 1.0, + device: Optional[torch.device] = None, + ): + super().__init__() + self.student_tokenizer = student_tokenizer + self.teacher_tokenizer_group = teacher_tokenizer_group + self.num_teachers = len(teacher_tokenizer_group) + + # Set teacher weights (default to equal weights) + if teacher_weights is None: + self.teacher_weights = [1.0 / self.num_teachers] * self.num_teachers + else: + if len(teacher_weights) != self.num_teachers: + raise ValueError( + f"Number of weights ({len(teacher_weights)}) must match number of teachers ({self.num_teachers})") + # Normalize weights to sum to 1 + weight_sum = sum(teacher_weights) + self.teacher_weights = [w / weight_sum for w in teacher_weights] + + self.max_length = max_length + self.beta = beta + self.gamma = gamma + self.temperature = temperature + + # Auto-detect NPU device if not specified + if device is None: + if torch.cuda.is_available(): + self.device = torch.device('cuda') + elif hasattr(torch, 'npu') and torch.npu.is_available(): + self.device = torch.device('npu:0') + print(f"Auto-detected NPU device: {self.device}") + else: + self.device = torch.device('cpu') + else: + self.device = device + + # Auto-select loss_type based on vocabulary coverage + if loss_type is None: + self.loss_type = self._auto_select_loss_type() + else: + self.loss_type = loss_type + + # Vocabulary sizes + self.student_vocab_size = len(student_tokenizer) + self.teacher_vocab_sizes = [len(tokenizer) for tokenizer in teacher_tokenizer_group] + + # Lazy initialization flags + self._projection_matrices_built = False + self.projection_matrices: list = [] + self.projection_student_indices_list: list = [] + self.projection_teacher_indices_list: list = [] + self._best_teacher_mappings: list = [] + + def _auto_select_loss_type(self) -> str: + # Calculate average coverage across all teachers + coverages = [] + for teacher_tokenizer in self.teacher_tokenizer_group: + coverage = self._calculate_vocab_coverage(self.student_tokenizer, teacher_tokenizer) + coverages.append(coverage) + + avg_coverage = sum(coverages) / len(coverages) + + # Select loss_type based on coverage threshold + if avg_coverage >= 0.7: # High coverage: use H-KL + return 'hkl' + else: # Low coverage: use P-KL + return 'pkl' + + def _calculate_vocab_coverage(self, student_tokenizer, teacher_tokenizer) -> float: + """ + Calculate vocabulary coverage between student and teacher tokenizers. + + Coverage = |intersection| / |union| + + Args: + student_tokenizer: Student model tokenizer + teacher_tokenizer: Teacher model tokenizer + + Returns: + Coverage ratio between 0 and 1 + """ + # Get vocabulary sets + student_vocab = set(student_tokenizer.get_vocab().keys()) + teacher_vocab = set(teacher_tokenizer.get_vocab().keys()) + + # Calculate intersection and union + intersection = student_vocab & teacher_vocab + union = student_vocab | teacher_vocab + + # Avoid division by zero + if len(union) == 0: + return 0.0 + + return len(intersection) / len(union) + + def _ensure_projection_matrices_built(self): + """Ensure projection matrices are built (lazy initialization with caching).""" + if self._projection_matrices_built: + return + # Use a lock to ensure only one thread builds the projection matrices + if not hasattr(self, '_build_lock'): + self._build_lock = threading.Lock() + + with self._build_lock: + # Check again inside the lock to avoid race condition + if self._projection_matrices_built: + return + + import time + start_time = time.perf_counter() + + # Generate cache key based on tokenizer configurations + cache_key = self._generate_cache_key() + + # Check if projection matrices are already cached + if cache_key in _PROJECTION_MATRIX_CACHE: + # Load from cache + cached_data = _PROJECTION_MATRIX_CACHE[cache_key] + self.projection_matrices = cached_data['projection_matrices'] + self.projection_student_indices_list = [t.to(self.device) if self.device is not None else t.clone() for + t in cached_data['projection_student_indices_list']] + self.projection_teacher_indices_list = [t.to(self.device) if self.device is not None else t.clone() for + t in cached_data['projection_teacher_indices_list']] + self.projection_values_list = [t.to(self.device) if self.device is not None else t.clone() for t in + cached_data['projection_values_list']] + else: + # Build projection matrices + self.projection_matrices = [] + self.projection_student_indices_list = [] + self.projection_teacher_indices_list = [] + self.projection_values_list = [] + for i, teacher_tokenizer in enumerate(self.teacher_tokenizer_group): + self._build_projection_matrix_for_teacher(teacher_tokenizer, i) + + # Cache the built matrices + _PROJECTION_MATRIX_CACHE[cache_key] = { + 'projection_matrices': self.projection_matrices, + 'projection_student_indices_list': self.projection_student_indices_list, + 'projection_teacher_indices_list': self.projection_teacher_indices_list, + 'projection_values_list': self.projection_values_list, + } + + # For H-KL: precompute the best teacher token for each student token for each teacher + if self.loss_type == 'hkl': + hkl_start_time = time.perf_counter() + self._best_teacher_mappings = [] + for i in range(self.num_teachers): + self._build_best_teacher_mapping_for_teacher(i) + + self._projection_matrices_built = True + + def _generate_cache_key(self) -> str: + """Generate a unique cache key based on tokenizer configurations.""" + # Create a hashable representation of the tokenizer configurations + config_data = { + 'student_vocab': self.student_tokenizer.get_vocab(), + 'teacher_vocabs': [tokenizer.get_vocab() for tokenizer in self.teacher_tokenizer_group], + 'max_length': self.max_length, + 'beta': self.beta, + 'gamma': self.gamma, + } + + # Use hash of the configuration data as cache key + config_bytes = pickle.dumps(config_data) + return hashlib.md5(config_bytes).hexdigest() + + def __call__( + self, + inputs, + outputs, + **kwargs, + ) -> LossOutput: + """Compute CTKD loss between student and multiple teacher models. + + Args: + inputs: Dict containing 'input_ids' and 'labels' for student model. + outputs: Dict containing 'logits' from student model. + teacher_logits_group: List of teacher model logits for each teacher. + teacher_topk_logprobs_group: List of teacher topk logprobs for each teacher. + teacher_topk_indices_group: List of teacher topk indices for each teacher. + **kwargs: Additional arguments. + + Returns: + LossOutput with the computed loss and number of tokens. + """ + # Ensure projection matrices are built (lazy initialization) + try: + self._ensure_projection_matrices_built() + except Exception as e: + import traceback + traceback.print_exc() + raise + # Extract student logits and labels + student_logits = outputs.get('logits') + if student_logits is None: + raise ValueError("Student logits not found in outputs") + + student_labels = inputs.get('labels') + if student_labels is None: + raise ValueError("Student labels not found in inputs") + + # Extract teacher logits group + teacher_logits_group = kwargs.get('teacher_logits_group') + if teacher_logits_group is None: + teacher_logits_group = outputs.get('teacher_logits_group') + teacher_topk_logprobs_group = kwargs.get('teacher_topk_logprobs_group') + teacher_topk_indices_group = kwargs.get('teacher_topk_indices_group') + # If we have topk format but not full logits, convert topk to logits for each teacher + if teacher_logits_group is None and teacher_topk_logprobs_group is not None and teacher_topk_indices_group is not None: + if len(teacher_topk_logprobs_group) != self.num_teachers or len( + teacher_topk_indices_group) != self.num_teachers: + raise ValueError( + f"Number of teachers in topk format ({len(teacher_topk_logprobs_group)}) must match number of teachers ({self.num_teachers})") + + teacher_logits_group = [] + for i in range(self.num_teachers): + teacher_topk_logprobs = teacher_topk_logprobs_group[i] + teacher_topk_indices = teacher_topk_indices_group[i] + + # Get vocabulary size for this teacher + vocab_size = self.teacher_vocab_sizes[i] + batch_size, seq_len, topk = teacher_topk_logprobs.shape + + # CRITICAL FIX: vLLM returns logprobs (log probabilities), not logits. + # We convert logprobs to probabilities using exp(). + # The resulting tensor represents teacher probabilities, not logits. + # These will be used directly as probabilities in _compute_pkl_loss/_compute_hkl_loss + # (NOT applying softmax again). + + # Create full probability tensor initialized with zeros + teacher_probs_full = torch.zeros( + (batch_size, seq_len, vocab_size), + dtype=teacher_topk_logprobs.dtype, + device=student_logits.device if student_logits is not None else teacher_topk_logprobs.device + ) + + # Convert logprobs to probabilities using exp() + teacher_topk_probs = torch.exp(teacher_topk_logprobs) + + # Scatter the topk probabilities into the full probability tensor + teacher_probs_full.scatter_( + dim=2, + index=teacher_topk_indices.to(teacher_probs_full.device), + src=teacher_topk_probs.to(teacher_probs_full.device) + ) + + # Store as logits_group for compatibility, but these are actually probabilities + # The _compute_pkl_loss/_compute_hkl_loss functions will use them directly + teacher_logits_group.append(teacher_probs_full) + + if teacher_logits_group is None: + raise ValueError("Teacher logits group not found in kwargs or outputs. " + "Provide either teacher_logits_group or (teacher_topk_logprobs_group + teacher_topk_indices_group)") + + if len(teacher_logits_group) != self.num_teachers: + raise ValueError( + f"Number of teacher logits ({len(teacher_logits_group)}) must match number of teachers ({self.num_teachers})") + + # Get labels - prefer teacher_labels from kwargs for CTKD loss mask + # The student labels may have -100 for most positions (e.g., only response tokens have valid labels), + # but for CTKD we want to compute loss over all positions where the teacher has valid predictions. + teacher_labels_group = kwargs.get('teacher_labels') + if teacher_labels_group is not None and len(teacher_labels_group) > 0: + labels = teacher_labels_group[0].to(student_logits.device) + else: + labels = inputs.get('labels') + if labels is None: + raise ValueError("labels not found in inputs") + + # Compute loss for each teacher and apply weighted average + total_loss = 0.0 + teacher_losses = [] + for i in range(self.num_teachers): + teacher_logits = teacher_logits_group[i] + weight = self.teacher_weights[i] + + if self.loss_type == 'pkl': + teacher_loss = self._compute_pkl_loss(student_logits, teacher_logits, labels, teacher_index=i) + elif self.loss_type == 'hkl': + teacher_loss = self._compute_hkl_loss(student_logits, teacher_logits, labels, teacher_index=i) + else: + raise ValueError(f"Unknown loss_type: {self.loss_type}. Use 'pkl' or 'hkl'") + + weighted_loss = weight * teacher_loss + total_loss += weighted_loss + teacher_losses.append({ + 'teacher_index': i, + 'loss_type': self.loss_type, + 'raw_loss': teacher_loss.item(), + 'weight': weight, + 'weighted_loss': weighted_loss.item() + }) + + # Print detailed loss information + print(f"\n=== CTKDLoss Detailed Breakdown ===") + print(f"Total Teachers: {self.num_teachers}") + print(f"Loss Type: {self.loss_type}") + print("-" * 50) + for loss_info in teacher_losses: + print(f"Teacher {loss_info['teacher_index']}:") + print(f" Loss Type: {loss_info['loss_type']}") + print(f" Raw Loss: {loss_info['raw_loss']:.6f}") + print(f" Weight: {loss_info['weight']:.4f}") + print(f" Weighted Loss: {loss_info['weighted_loss']:.6f}") + print("-" * 50) + print(f"Total Loss: {total_loss.item():.6f}") + print("=" * 50) + + # Use teacher_labels for num_tokens calculation if available + if teacher_labels_group is not None and len(teacher_labels_group) > 0: + num_tokens = teacher_labels_group[0].ne(-100).sum().item() + else: + num_tokens = labels.ne(-100).sum().item() if labels is not None else 0 + loss_output = LossOutput(loss=total_loss, num_tokens=num_tokens) + return loss_output + + def _build_projection_matrix_for_teacher(self, teacher_tokenizer, teacher_index): + """ + Build the sparse projection matrix W for a specific teacher tokenizer. + + The construction follows the X-Token paper algorithm: + + Step 1: Initialize W[s,t] = 0 for all s in V_S, t in V_T + + Step 2: Exact match - For each student token s: + - Decode s to text + - If text matches a teacher token t's decoded text: + W[s, t] = 1 + + Step 3: Multi-token decoding match - For unmatched student tokens s: + - Decode s to text + - Encode text with teacher tokenizer -> (t[0], ..., t[ℓ-1]) + - If ℓ < L (max_length): + For i in [0, ℓ-1]: + W[s, t[i]] = β * γ^i + + The resulting matrix maps student token probabilities to teacher token space, + enabling KL divergence computation across different vocabularies. + + Note: This implementation directly builds sparse COO format to avoid + memory issues with large vocabulary sizes. + """ + teacher_vocab_size = len(teacher_tokenizer) + + # Check memory requirements for dense matrix (for reference only) + matrix_size_gb = (self.student_vocab_size * teacher_vocab_size * 4) / (1024 ** 3) + + # Use lists to store sparse matrix entries (COO format) + # This avoids creating a huge dense matrix + student_indices = [] + teacher_indices = [] + values = [] + + # Get vocabulary mappings + student_vocab = self.student_tokenizer.get_vocab() # {token_str: token_id} + teacher_vocab = teacher_tokenizer.get_vocab() + + # Track which student tokens have been matched + matched_student_ids = set() + + # Step 2: Exact match - find tokens with identical text representation + # Build a mapping from token text to teacher token id for efficient lookup + teacher_token_text_to_id = {} + for token_id in range(teacher_vocab_size): + # Decode each teacher token to its text representation + # skip_special_tokens=False to preserve special tokens + token_text = teacher_tokenizer.decode( + [token_id], + skip_special_tokens=False + ).strip() + teacher_token_text_to_id[token_text] = token_id + + # Match student tokens to teacher tokens by text + for student_id in range(self.student_vocab_size): + # Decode student token to text + student_token_text = self.student_tokenizer.decode( + [student_id], + skip_special_tokens=False + ).strip() + + # Check if this text exists in teacher vocabulary + if student_token_text in teacher_token_text_to_id: + teacher_id = teacher_token_text_to_id[student_token_text] + # Store in sparse format directly + student_indices.append(student_id) + teacher_indices.append(teacher_id) + values.append(1.0) + matched_student_ids.add(student_id) + + # Step 3: Multi-token decoding match for unmatched student tokens + # For each unmatched student token, decode to text and re-encode with teacher tokenizer + unmatched_count = 0 + for student_id in range(self.student_vocab_size): + if student_id in matched_student_ids: + continue # Skip already matched tokens + + # Decode student token to raw text + text = self.student_tokenizer.decode( + [student_id], + skip_special_tokens=False + ) + + # Skip empty text + if not text or not text.strip(): + continue + + # Encode with teacher tokenizer + teacher_token_ids = teacher_tokenizer.encode( + text, + add_special_tokens=False + ) + + # Get the length of encoded sequence + seq_length = len(teacher_token_ids) + + # Only assign weights if sequence length < max_length (L) + if seq_length > 0 and seq_length < self.max_length: + for i, teacher_token_id in enumerate(teacher_token_ids): + # Weight follows exponential decay: β * γ^i + # Earlier tokens get higher weights + weight = self.beta * (self.gamma ** i) + # Store in sparse format directly + student_indices.append(student_id) + teacher_indices.append(teacher_token_id) + values.append(weight) + unmatched_count += 1 + + # Convert to tensors directly on target device to avoid device transfers + student_indices_tensor = torch.tensor(student_indices, dtype=torch.long, device=self.device) + teacher_indices_tensor = torch.tensor(teacher_indices, dtype=torch.long, device=self.device) + values_tensor = torch.tensor(values, dtype=torch.float32, + device=self.device) # Use float32 for better gradient computation + + # Store as separate tensors (COO sparse format) + self.projection_student_indices_list.append(student_indices_tensor) + self.projection_teacher_indices_list.append(teacher_indices_tensor) + self.projection_values_list.append(values_tensor) + + # Store None for dense matrix to save memory + self.projection_matrices.append(None) + + # Calculate actual memory usage + sparse_memory_mb = (student_indices_tensor.numel() * 8 + + teacher_indices_tensor.numel() * 8 + + values_tensor.numel() * 4) / (1024 ** 2) + + print( + f"Built projection matrix for teacher {teacher_index}: {len(student_indices)} mappings, memory: {sparse_memory_mb:.2f} MB") + + def _build_best_teacher_mapping_for_teacher(self, teacher_index): + """ + Build the best teacher token mapping for H-KL loss for a specific teacher. + + For each student token, find the teacher token with the highest projection weight: + t* = argmax_{t' in V_T} W[s, t'] + constraint: W[s, t*] > 0 + + This creates a one-to-one mapping for heuristic KL divergence computation. + + Note: This implementation uses sparse COO format to avoid memory issues. + """ + if teacher_index >= len(self.projection_student_indices_list): + raise ValueError(f"Projection matrix for teacher {teacher_index} must be built first") + + # Use sparse COO format data + student_indices = self.projection_student_indices_list[teacher_index] + teacher_indices = self.projection_teacher_indices_list[teacher_index] + values = self.projection_values_list[teacher_index].float() + + # Initialize mapping with -1 (no mapping) + best_teacher_mapping = torch.full( + (self.student_vocab_size,), + -1, + dtype=torch.long + ) + max_weights = torch.zeros(self.student_vocab_size, dtype=torch.float32) + + # Find the best teacher token for each student token + # Process in chunks to avoid memory issues + chunk_size = 100000 + for i in range(0, len(student_indices), chunk_size): + chunk_student = student_indices[i:i + chunk_size] + chunk_teacher = teacher_indices[i:i + chunk_size] + chunk_values = values[i:i + chunk_size] + + # For each entry, check if it's the best weight for that student token + for j in range(len(chunk_student)): + s_id = chunk_student[j].item() + t_id = chunk_teacher[j].item() + w = chunk_values[j].item() + + if w > max_weights[s_id]: + max_weights[s_id] = w + best_teacher_mapping[s_id] = t_id + + self._best_teacher_mappings.append(best_teacher_mapping) + + def _compute_pkl_loss( + self, + student_logits: torch.Tensor, + teacher_probs: torch.Tensor, # NOTE: These are already probabilities (from exp(logprobs)), NOT logits + labels: torch.Tensor, + teacher_index: int = 0, + ) -> torch.Tensor: + """ + Compute P-KL (Partition-free KL) loss. + + P-KL projects the student distribution to teacher vocabulary space using + the full projection matrix W, then computes KL divergence: + + p̃_S[t] = Σ_{s in V_S} W[s,t] * p_S[s] + L_P = KL(p_T || p̃_S) + + CRITICAL: The `teacher_probs` input is expected to be already in probability space + (converted from vLLM logprobs via exp()). NO softmax should be applied to it. + + Args: + student_logits: [batch, seq_len, student_vocab_size] - raw logits from student model + teacher_probs: [batch, seq_len, teacher_vocab_size] - probabilities from teacher model + labels: [batch, seq_len] + teacher_index: Index of the teacher model + + Returns: + Scalar loss value. + """ + # Shift logits and labels for next-token prediction + shift_student_logits = student_logits[..., :-1, :].contiguous() + shift_teacher_probs = teacher_probs[..., :-1, :].contiguous().to(student_logits.device) + shift_labels = labels[..., 1:].contiguous().to(student_logits.device) + + # Create loss mask + loss_mask = (shift_labels != -100).float() + + # Compute student probabilities with temperature + student_probs = F.softmax(shift_student_logits / self.temperature, dim=-1) + + # ============ DEBUG LOGS ============ + print(f"[CTKD DEBUG] _compute_pkl_loss: student_logits shape: {student_logits.shape}") + print(f"[CTKD DEBUG] _compute_pkl_loss: teacher_probs shape: {teacher_probs.shape}") + print(f"[CTKD DEBUG] _compute_pkl_loss: labels shape: {labels.shape}") + print(f"[CTKD DEBUG] _compute_pkl_loss: shift_student_logits shape: {shift_student_logits.shape}") + print(f"[CTKD DEBUG] _compute_pkl_loss: shift_teacher_probs shape: {shift_teacher_probs.shape}") + print(f"[CTKD DEBUG] _compute_pkl_loss: shift_labels shape: {shift_labels.shape}") + print(f"[CTKD DEBUG] _compute_pkl_loss: loss_mask sum: {loss_mask.sum().item()}, total: {loss_mask.numel()}") + print( + f"[CTKD DEBUG] _compute_pkl_loss: student_logits range: [{shift_student_logits.min().item():.4f}, {shift_student_logits.max().item():.4f}]") + print( + f"[CTKD DEBUG] _compute_pkl_loss: student_probs range: [{student_probs.min().item():.6f}, {student_probs.max().item():.6f}]") + print( + f"[CTKD DEBUG] _compute_pkl_loss: student_probs sum per position: [{student_probs.sum(dim=-1).min().item():.4f}, {student_probs.sum(dim=-1).max().item():.4f}]") + print( + f"[CTKD DEBUG] _compute_pkl_loss: teacher_probs range: [{shift_teacher_probs.min().item():.6f}, {shift_teacher_probs.max().item():.6f}]") + print( + f"[CTKD DEBUG] _compute_pkl_loss: teacher_probs sum per position: [{shift_teacher_probs.sum(dim=-1).min().item():.6f}, {shift_teacher_probs.sum(dim=-1).max().item():.6f}]") + print( + f"[CTKD DEBUG] _compute_pkl_loss: teacher_nonzero_count per position: [{(shift_teacher_probs > 0).sum(dim=-1).min().item()}, {(shift_teacher_probs > 0).sum(dim=-1).max().item()}]") + print(f"[CTKD DEBUG] _compute_pkl_loss: temperature: {self.temperature}") + # ==================================== + + # CRITICAL: teacher_probs are already probabilities (converted from vLLM logprobs via exp()). + # They are a top-k truncated distribution - most entries are 0. + # We should NOT apply softmax or temperature scaling to them. + # The teacher distribution is used as-is. + + # Align sequence lengths by taking the minimum + student_seq_len = student_probs.shape[1] + teacher_seq_len = shift_teacher_probs.shape[1] + min_seq_len = min(student_seq_len, teacher_seq_len) + + # Use the same sequence length for both student and teacher probabilities + student_probs = student_probs[:, :min_seq_len, :] + shift_teacher_probs = shift_teacher_probs[:, :min_seq_len, :] + loss_mask = loss_mask[:, :min_seq_len] + + # Project student probabilities to teacher vocabulary space + # p̃_S[t] = Σ_s W[s,t] * p_S[s] + # Instead of using sparse matrix multiplication (not supported on NPU), + # we use scatter_add with stored indices and values + + # Use actual teacher vocabulary size from teacher_probs tensor + _, _, teacher_vocab_size = shift_teacher_probs.shape + + # Move projection indices and values to device for the specific teacher + student_indices = self.projection_student_indices_list[teacher_index].to(student_probs.device) + teacher_indices = self.projection_teacher_indices_list[teacher_index].to(student_probs.device) + proj_values = self.projection_values_list[teacher_index].to(student_probs.device).float() + + # Filter teacher indices to ensure they are within the valid range + # This handles cases where the actual teacher vocab size differs from expected + valid_mask = teacher_indices < teacher_vocab_size + if not valid_mask.all(): + # Some teacher indices are out of bounds, filter them + student_indices = student_indices[valid_mask] + teacher_indices = teacher_indices[valid_mask] + proj_values = proj_values[valid_mask] + + # Get student probabilities for the non-zero projection entries + # student_probs: [batch, seq_len, student_vocab] + # student_indices: [num_non_zero] + # We need to gather: student_probs[:, :, student_indices] + selected_student_probs = student_probs.index_select( + dim=-1, + index=student_indices + ) # [batch, seq_len, num_non_zero] + + # Multiply by projection weights + weighted_probs = selected_student_probs * proj_values.unsqueeze(0).unsqueeze(0) + # [batch, seq_len, num_non_zero] + + # Scatter add to teacher vocabulary space + projected_student_probs = torch.zeros( + batch_size, seq_len, teacher_vocab_size, + device=student_probs.device, + dtype=student_probs.dtype + ) + + # Expand teacher_indices for scatter_add + # teacher_indices: [num_non_zero] -> [batch, seq_len, num_non_zero] + expanded_teacher_indices = teacher_indices.unsqueeze(0).unsqueeze(0).expand( + batch_size, seq_len, -1 + ) + + # Scatter add: accumulate weighted probabilities to teacher tokens + # Ensure indices and source are on the same device as target + projected_student_probs.scatter_add_( + dim=2, + index=expanded_teacher_indices.to(projected_student_probs.device), + src=weighted_probs.to(projected_student_probs.device) + ) + + # CRITICAL FIX: Do NOT renormalize projected_student_probs. + # The projection matrix W[s,t] is designed so that Σ_t W[s,t] = 1 for each student token s, + # and the scatter_add operation preserves the total probability mass. + # Renormalization would destroy the probability structure and inflate the loss. + # The KL divergence will only be computed over positions where teacher has non-zero probability, + # so the missing probability mass in teacher's top-k truncated distribution is handled correctly. + + # Debug: Check dimensions before KL divergence + if projected_student_probs.size(-1) != shift_teacher_probs.size(-1): + raise ValueError( + f"Vocabulary dimension mismatch: projected_student_probs has size {projected_student_probs.size(-1)}, " + f"teacher_probs has size {shift_teacher_probs.size(-1)}. " + f"Projection matrix was built for teacher vocab size {self.teacher_vocab_size}, " + f"but actual teacher model has vocab size {shift_teacher_probs.size(-1)}. " + f"This suggests the teacher model used at runtime has a different vocabulary than the teacher tokenizer used for projection matrix construction." + ) + teacher_nonzero_mask = shift_teacher_probs > 0 # [batch, seq_len, teacher_vocab] + + # Compute KL divergence only over teacher non-zero positions + # KL(P||Q) = Σ P(x) * log(P(x)/Q(x)) + # We compute this manually to handle the masking correctly + log_projected_student = torch.log(projected_student_probs + 1e-8) + + # Element-wise KL contribution: P * log(P/Q) + # Only compute for positions where teacher has non-zero probability + kl_contrib = shift_teacher_probs * ( + torch.log(shift_teacher_probs + 1e-8) - log_projected_student + ) + + # Zero out contributions from teacher-zero positions + kl_contrib = kl_contrib * teacher_nonzero_mask + + # Sum over vocabulary dimension + kl_div = kl_contrib.sum(dim=-1) # [batch, seq_len] + + # Apply mask and average + masked_kl = kl_div * loss_mask + loss = masked_kl.sum() / (loss_mask.sum() + 1e-8) + + # ============ DEBUG LOGS: KL Computation ============ + print( + f"[CTKD DEBUG] _compute_pkl_loss: projected_student_probs range: [{projected_student_probs.min().item():.6f}, {projected_student_probs.max().item():.6f}]") + print( + f"[CTKD DEBUG] _compute_pkl_loss: projected_student_probs sum per position: [{projected_student_probs.sum(dim=-1).min().item():.6f}, {projected_student_probs.sum(dim=-1).max().item():.6f}]") + print( + f"[CTKD DEBUG] _compute_pkl_loss: log_projected_student range: [{log_projected_student.min().item():.4f}, {log_projected_student.max().item():.4f}]") + print( + f"[CTKD DEBUG] _compute_pkl_loss: kl_contrib range: [{kl_contrib.min().item():.6f}, {kl_contrib.max().item():.6f}]") + print( + f"[CTKD DEBUG] _compute_pkl_loss: kl_div range (per position): [{kl_div.min().item():.4f}, {kl_div.max().item():.4f}]") + print(f"[CTKD DEBUG] _compute_pkl_loss: kl_div mean (unmasked): {kl_div.mean().item():.4f}") + print( + f"[CTKD DEBUG] _compute_pkl_loss: masked_kl range: [{masked_kl.min().item():.4f}, {masked_kl.max().item():.4f}]") + print(f"[CTKD DEBUG] _compute_pkl_loss: final loss: {loss.item():.6f}") + + # Check for potential issues + if projected_student_probs.max().item() < 1e-6: + print(f"[CTKD WARNING] projected_student_probs is nearly zero! This indicates projection matrix issues.") + if kl_div.max().item() > 50: + print(f"[CTKD WARNING] Very high KL divergence detected! Max kl_div: {kl_div.max().item():.4f}") + # ==================================== + + return loss + + def _compute_hkl_loss( + self, + student_logits: torch.Tensor, + teacher_probs: torch.Tensor, # NOTE: These are already probabilities (from exp(logprobs)), NOT logits + labels: torch.Tensor, + teacher_index: int = 0, + ) -> torch.Tensor: + # Shift logits and labels for next-token prediction + shift_student_logits = student_logits[..., :-1, :].contiguous() + shift_teacher_probs = teacher_probs[..., :-1, :].contiguous().to(student_logits.device) + shift_labels = labels[..., 1:].contiguous().to(student_logits.device) + + # Create loss mask + loss_mask = (shift_labels != -100).float() + + # Compute student probabilities with temperature + student_probs = F.softmax(shift_student_logits / self.temperature, dim=-1) + + # ============ DEBUG LOGS ============ + print(f"[CTKD DEBUG] _compute_hkl_loss: student_logits shape: {student_logits.shape}") + print(f"[CTKD DEBUG] _compute_hkl_loss: teacher_probs shape: {teacher_probs.shape}") + print(f"[CTKD DEBUG] _compute_hkl_loss: labels shape: {labels.shape}") + print(f"[CTKD DEBUG] _compute_hkl_loss: shift_student_logits shape: {shift_student_logits.shape}") + print(f"[CTKD DEBUG] _compute_hkl_loss: shift_teacher_probs shape: {shift_teacher_probs.shape}") + print(f"[CTKD DEBUG] _compute_hkl_loss: shift_labels shape: {shift_labels.shape}") + print(f"[CTKD DEBUG] _compute_hkl_loss: loss_mask sum: {loss_mask.sum().item()}, total: {loss_mask.numel()}") + print( + f"[CTKD DEBUG] _compute_hkl_loss: student_logits range: [{shift_student_logits.min().item():.4f}, {shift_student_logits.max().item():.4f}]") + print( + f"[CTKD DEBUG] _compute_hkl_loss: student_probs range: [{student_probs.min().item():.6f}, {student_probs.max().item():.6f}]") + print( + f"[CTKD DEBUG] _compute_hkl_loss: student_probs sum per position: [{student_probs.sum(dim=-1).min().item():.4f}, {student_probs.sum(dim=-1).max().item():.4f}]") + print( + f"[CTKD DEBUG] _compute_hkl_loss: teacher_probs range: [{shift_teacher_probs.min().item():.6f}, {shift_teacher_probs.max().item():.6f}]") + print( + f"[CTKD DEBUG] _compute_hkl_loss: teacher_probs sum per position: [{shift_teacher_probs.sum(dim=-1).min().item():.6f}, {shift_teacher_probs.sum(dim=-1).max().item():.6f}]") + print( + f"[CTKD DEBUG] _compute_hkl_loss: teacher_nonzero_count per position: [{(shift_teacher_probs > 0).sum(dim=-1).min().item()}, {(shift_teacher_probs > 0).sum(dim=-1).max().item()}]") + print(f"[CTKD DEBUG] _compute_hkl_loss: temperature: {self.temperature}") + # ==================================== + student_seq_len = student_probs.shape[1] + teacher_seq_len = shift_teacher_probs.shape[1] + min_seq_len = min(student_seq_len, teacher_seq_len) + + # Use the same sequence length for both student and teacher probabilities + student_probs = student_probs[:, :min_seq_len, :] + shift_teacher_probs = shift_teacher_probs[:, :min_seq_len, :] + loss_mask = loss_mask[:, :min_seq_len] + + # Get best teacher token for each student token for the specific teacher + best_mapping = self._best_teacher_mappings[teacher_index].to(student_probs.device) + + # Select student probabilities for mapped teacher tokens + # student_probs: [batch, seq_len, student_vocab] + # best_mapping: [student_vocab] -> teacher token id for each student token + # We need to gather the mapped probabilities + batch_size, seq_len, student_vocab = student_probs.shape + teacher_vocab = shift_teacher_probs.shape[-1] + + # Create output tensor for mapped student probabilities + mapped_student_probs = torch.zeros( + batch_size, seq_len, teacher_vocab, + device=student_probs.device, + dtype=student_probs.dtype + ) + + # For each student token with a valid mapping, add its probability to the mapped teacher token + valid_mask = best_mapping >= 0 # [student_vocab] + valid_student_ids = torch.where(valid_mask)[0] + valid_teacher_ids = best_mapping[valid_student_ids] + # Process in chunks to avoid memory issues with large vocabularies + chunk_size = 10000 # Process 10k tokens at a time + for i in range(0, len(valid_student_ids), chunk_size): + chunk_student_ids = valid_student_ids[i:i + chunk_size] + chunk_teacher_ids = valid_teacher_ids[i:i + chunk_size] + + # Scatter student probabilities to teacher vocabulary positions + # Ensure indices and source are on the same device as target + mapped_student_probs.scatter_add_( + dim=2, + index=chunk_teacher_ids.unsqueeze(0).unsqueeze(0).expand(batch_size, seq_len, -1).to( + mapped_student_probs.device), + src=student_probs[:, :, chunk_student_ids].to(mapped_student_probs.device) + ) + + # CRITICAL FIX: Do NOT renormalize mapped_student_probs. + # The mapping preserves probability mass only for tokens with valid mappings. + # Renormalization would destroy the probability structure. + + # ============ DEBUG LOGS: Mapping Statistics ============ + print( + f"[CTKD DEBUG] _compute_hkl_loss: best_mapping valid count: {valid_mask.sum().item()} / {best_mapping.numel()}") + print(f"[CTKD DEBUG] _compute_hkl_loss: valid_student_ids count: {len(valid_student_ids)}") + print( + f"[CTKD DEBUG] _compute_hkl_loss: mapped_student_probs range: [{mapped_student_probs.min().item():.6f}, {mapped_student_probs.max().item():.6f}]") + print( + f"[CTKD DEBUG] _compute_hkl_loss: mapped_student_probs sum per position: [{mapped_student_probs.sum(dim=-1).min().item():.6f}, {mapped_student_probs.sum(dim=-1).max().item():.6f}]") + # ==================================== + + # CRITICAL FIX: Compute KL divergence only over positions where teacher has non-zero probability. + teacher_nonzero_mask = shift_teacher_probs > 0 # [batch, seq_len, teacher_vocab] + + # Compute KL divergence: KL(p_T || mapped_p_S) + # Only compute for positions where teacher has non-zero probability + log_mapped_student = torch.log(mapped_student_probs + 1e-8) + + # Element-wise KL contribution: P * log(P/Q) + kl_contrib = shift_teacher_probs * ( + torch.log(shift_teacher_probs + 1e-8) - log_mapped_student + ) + # Zero out contributions from teacher-zero positions + kl_contrib = kl_contrib * teacher_nonzero_mask + + # Sum over vocabulary dimension + kl_div = kl_contrib.sum(dim=-1) + + # Apply mask and average + masked_kl = kl_div * loss_mask + loss = masked_kl.sum() / (loss_mask.sum() + 1e-8) + + # ============ DEBUG LOGS: KL Computation ============ + print( + f"[CTKD DEBUG] _compute_hkl_loss: log_mapped_student range: [{log_mapped_student.min().item():.4f}, {log_mapped_student.max().item():.4f}]") + print( + f"[CTKD DEBUG] _compute_hkl_loss: kl_contrib range: [{kl_contrib.min().item():.6f}, {kl_contrib.max().item():.6f}]") + print( + f"[CTKD DEBUG] _compute_hkl_loss: kl_div range (per position): [{kl_div.min().item():.4f}, {kl_div.max().item():.4f}]") + print(f"[CTKD DEBUG] _compute_hkl_loss: kl_div mean (unmasked): {kl_div.mean().item():.4f}") + print( + f"[CTKD DEBUG] _compute_hkl_loss: masked_kl range: [{masked_kl.min().item():.4f}, {masked_kl.max().item():.4f}]") + print(f"[CTKD DEBUG] _compute_hkl_loss: final loss: {loss.item():.6f}") + + # Check for potential issues + if mapped_student_probs.max().item() < 1e-6: + print(f"[CTKD WARNING] mapped_student_probs is nearly zero! This indicates mapping issues.") + if kl_div.max().item() > 50: + print(f"[CTKD WARNING] Very high KL divergence detected! Max kl_div: {kl_div.max().item():.4f}") + # ==================================== + + return loss + + def get_projection_matrix(self, teacher_index: int = 0) -> Optional[torch.Tensor]: + """Return the projection matrix W for a specific teacher. + + Args: + teacher_index: Index of the teacher model (default: 0). + + Returns: + Tensor of shape [student_vocab_size, teacher_vocab_size] or None if using sparse format. + Note: For large vocabularies, this returns None to save memory. + Use get_sparse_projection_data() to get the sparse representation. + """ + # Ensure projection matrices are built before accessing + self._ensure_projection_matrices_built() + + if teacher_index >= len(self.projection_matrices): + raise ValueError(f"Projection matrix for teacher {teacher_index} has not been built") + return self.projection_matrices[teacher_index] + + def get_sparse_projection_data(self, teacher_index: int = 0) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return the sparse projection matrix data in COO format. + + Args: + teacher_index: Index of the teacher model (default: 0). + + Returns: + Tuple of (student_indices, teacher_indices, values) representing the sparse matrix. + """ + # Ensure projection matrices are built before accessing + self._ensure_projection_matrices_built() + + if teacher_index >= len(self.projection_student_indices_list): + raise ValueError(f"Projection matrix for teacher {teacher_index} has not been built") + + return ( + self.projection_student_indices_list[teacher_index], + self.projection_teacher_indices_list[teacher_index], + self.projection_values_list[teacher_index], + ) + + def get_mapping_statistics(self, teacher_index: int = 0) -> Dict: + """Return statistics about the projection matrix for a specific teacher. + + Args: + teacher_index: Index of the teacher model (default: 0). + + Returns: + Dict containing: + - total_student_tokens: Total number of student tokens + - exact_matched: Number of tokens with exact match (W[s,t]=1) + - multi_token_matched: Number of tokens with multi-token mapping + - unmatched: Number of tokens without any mapping + - sparsity: Fraction of zero elements in the matrix + """ + # Ensure projection matrices are built before accessing statistics + self._ensure_projection_matrices_built() + + if teacher_index >= len(self.projection_student_indices_list): + raise ValueError(f"Projection matrix for teacher {teacher_index} has not been built") + + # Use the stored indices and values (COO format) for the specific teacher + student_indices = self.projection_student_indices_list[teacher_index] + values = self.projection_values_list[teacher_index].float() # Convert from half to float for comparison + + # Total number of non-zero elements + nnz = student_indices.numel() + total_elements = self.student_vocab_size * self.teacher_vocab_sizes[teacher_index] + + # Sparsity = fraction of zero elements + sparsity = 1.0 - (nnz / total_elements) + + # Count exact matches (weight == 1) + exact_match_mask = (values == 1.0) + exact_matched_students = student_indices[exact_match_mask].unique() + exact_matched = exact_matched_students.numel() + + # Count multi-token matches (0 < weight < 1) + multi_token_mask = (values > 0) & (values < 1.0) + multi_token_students = student_indices[multi_token_mask].unique() + # Exclude students that already have exact matches + multi_token_students = multi_token_students[ + ~multi_token_students.unsqueeze(1).eq(exact_matched_students.unsqueeze(0)).any(dim=1) + ] + multi_token_matched = multi_token_students.numel() + + # Count unmatched: total - exact_matched - multi_token_matched + unmatched = self.student_vocab_size - exact_matched - multi_token_matched + + return { + 'total_student_tokens': self.student_vocab_size, + 'exact_matched': exact_matched, + 'multi_token_matched': multi_token_matched, + 'unmatched': unmatched, + 'sparsity': sparsity + } \ No newline at end of file