diff --git a/cookbook/rl/mopd/ctkd_npu.py b/cookbook/rl/mopd/ctkd_npu.py new file mode 100644 index 00000000..66adb4d6 --- /dev/null +++ b/cookbook/rl/mopd/ctkd_npu.py @@ -0,0 +1,352 @@ +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', 'ms://Qwen/qwen3.5-2B') +TEACHER_MODEL_ID = os.environ.get('TEACHER_MODEL_ID', 'ms://Qwen/Qwen3-0.6B') +TEACHER_MODEL_ID_2 = os.environ.get('TEACHER_MODEL_ID_2', 'ms://Qwen/qwen2.5-0.5b-instruct') +DATASET_ID = os.environ.get('DATASET_ID', 'ms://hjh0119/shareAI-Llama3-DPO-zh-en-emoji') + +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 = MODEL_GPUS + (SAMPLER_GPUS if SHARED_TEACHER_GPUS else SAMPLER_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-4)) +GRADIENT_ACCUMULATION_STEPS = int(os.environ.get('GRADIENT_ACCUMULATION_STEPS', 2)) + +CTKD_TEMPERATURE = float(os.environ.get('CTKD_TEMPERATURE', 1.0)) +CTKD_MAX_LENGTH = int(os.environ.get('CTKD_MAX_LENGTH', 4)) +CTKD_BETA = float(os.environ.get('CTKD_BETA', 0.9)) +CTKD_GAMMA = float(os.environ.get('CTKD_GAMMA', 0.1)) +CTKD_LOSS_TYPE = os.environ.get('CTKD_LOSS_TYPE', 'pkl') # 'pkl' or 'hkl' +CTKD_TOPK = int(os.environ.get('CTKD_TOPK', 64)) +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: + """Convert vLLM topk_prompt_logprobs to CTKDLoss teacher_output format. + + Args: + topk_prompt_logprobs_batch: List of per-input topk_prompt_logprobs. + Each is List[Optional[List[(token_id, logprob)]]] of shape [seq_len, topk]. + topk: Number of top-k logits to extract. + + Returns: + Dict with 'teacher_topk_logprobs' [batch, seq_len, topk] and + 'teacher_topk_indices' [batch, seq_len, topk] tensors. + """ + 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: + # First position is None, fill with placeholder + 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) + + # Pad to same seq_len within batch + 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), + } + + +# ── Dataset ─────────────────────────────────────────────────────────────────── + +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 +def train(): + """Main training function for MOPD with CTKDLoss.""" + # Initialize device groups based on shared mode + if SHARED_TEACHER_GPUS: + # Shared mode: both teacher samplers use the same GPU resources + device_groups = [ + DeviceGroup(name='student_model', ranks=MODEL_GPUS, device_type='npu'), + DeviceGroup(name='teacher_sampler', ranks=SAMPLER_GPUS, device_type='npu'), + ] + else: + # Separate mode: each teacher sampler gets its own GPU resources + device_groups = [ + DeviceGroup(name='student_model', ranks=MODEL_GPUS, device_type='npu'), + DeviceGroup(name='teacher_sampler', ranks=SAMPLER_GPUS, device_type='npu'), + DeviceGroup(name='teacher_sampler_2', 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, + ) + + # ── Student model (trainable) ────────────────────────────────────────────── + student_model = TransformersModel( + model_id=STUDENT_MODEL_ID, + device_mesh=model_mesh, + remote_group='student_model', + ) + + # 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 ───────────────────────────────────────────────────── + # Get tokenizers for cross-tokenizer alignment + 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) + teacher_tokenizer_2 = AutoTokenizer.from_pretrained(TEACHER_MODEL_ID_2, trust_remote_code=True) + + loss_fn = CTKDLoss( + student_tokenizer=student_tokenizer, + teacher_tokenizer_group=[teacher_tokenizer, teacher_tokenizer_2], + max_length=CTKD_MAX_LENGTH, + beta=CTKD_BETA, + gamma=CTKD_GAMMA, + loss_type=CTKD_LOSS_TYPE, + temperature=CTKD_TEMPERATURE, + ) + student_model.set_loss(loss_fn, adapter_name=ADAPTER_NAME) + student_model.set_template('QwenTemplate', model_id=STUDENT_MODEL_ID, adapter_name=ADAPTER_NAME) + + # 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 + stats = loss_fn.get_mapping_statistics() + logger.info(f'CTKD Projection Matrix Statistics: {stats}') + + #── Teacher vLLM samplers (for logits) ───────────────────────────────────── + if SHARED_TEACHER_GPUS: + # Shared mode: use different instance_id to avoid naming conflicts + 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) + + teacher_sampler_2 = vLLMSampler( + model_id=TEACHER_MODEL_ID_2, + 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', # Same remote_group but different instance_id + instance_id='teacher_2' + ) + teacher_sampler_2.set_template('QwenTemplate', model_id=TEACHER_MODEL_ID_2) + else: + # Separate mode: each teacher sampler has its own GPU resources + 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) + + teacher_sampler_2 = vLLMSampler( + model_id=TEACHER_MODEL_ID_2, + 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_2', + ) + teacher_sampler_2.set_template('QwenTemplate', model_id=TEACHER_MODEL_ID_2) + + # ── DataLoader (full-text: prompt + response) ────────────────────────────── + # Dataset is pre-encoded with student tokenizer + # Teacher will decode text and re-encode with its own tokenizer + dataloader = DataLoader( + dataset=create_dataset(), # 调用函数获取数据集对象 + batch_size=BATCH_SIZE, + min_batch_size=BATCH_SIZE, + device_mesh=model_mesh, + remote_group='student_model', + ) + + logger.info(get_device_placement()) + logger.info(f'MOPD CTKD Training | student={STUDENT_MODEL_ID} teachers={TEACHER_MODEL_ID}, {TEACHER_MODEL_ID_2}') + logger.info(f' T={CTKD_TEMPERATURE} loss_type={CTKD_LOSS_TYPE} topk={CTKD_TOPK}') + logger.info(f' batch_size={BATCH_SIZE} lr={LEARNING_RATE} max_steps={MAX_STEPS}') + + # ── Training Loop ────────────────────────────────────────────────────────── + optim_step = 0 + for batch in dataloader: + if optim_step >= MAX_STEPS: + break + # 在 Ray 模式下，batch 是一个可调用对象，需要调用它来获取实际数据 + if callable(batch): + batch = batch() + + # 1. Decode student-encoded tokens back to text for teacher + # batch contains input_ids encoded with student tokenizer + # We need to decode them and re-encode with teacher tokenizer + from twinkle.data_format import Trajectory + + teacher_inputs = [] + for item in batch: + # Decode student tokens back to text, then re-encode with teacher tokenizer + text = student_tokenizer.decode(item['input_ids'], skip_special_tokens=False) + teacher_encoded = teacher_tokenizer.encode(text, add_special_tokens=False) + teacher_inputs.append({ + 'input_ids': teacher_encoded, + 'labels': item.get('labels', [-100] * len(teacher_encoded)), + }) + + # 2. Teachers compute top-k logprobs on the full sequences + # max_tokens=0: don't generate new content, just compute logits on input + teacher_response = teacher_sampler.sample( + teacher_inputs, # Trajectory format - teacher will encode with its tokenizer + SamplingParams(max_tokens=1, temperature=1.0, prompt_logprobs=CTKD_TOPK), # max_tokens=1 for logprobs only + ) + + teacher_response_2 = teacher_sampler_2.sample( + teacher_inputs, # Same input for second teacher + SamplingParams(max_tokens=1, temperature=1.0, prompt_logprobs=CTKD_TOPK), + ) + + # 3. Convert teacher responses to input format (already encoded with teacher tokenizers) + teacher_input_data = [seq.new_input_feature for resp in teacher_response for seq in resp.sequences] + teacher_input_data_2 = [seq.new_input_feature for resp in teacher_response_2 for seq in resp.sequences] + + # 4. Prepare teacher output for CTKDLoss (for both teachers) + topk_data = convert_topk_prompt_logprobs( + [resp.topk_prompt_logprobs for resp in teacher_response], + topk=CTKD_TOPK, + ) + + topk_data_2 = convert_topk_prompt_logprobs( + [resp.topk_prompt_logprobs for resp in teacher_response_2], + topk=CTKD_TOPK, + ) + + # Get teacher input_ids and labels from teacher_input_data + import torch.nn.utils.rnn as rnn_utils + teacher_input_ids_list = [torch.tensor(item['input_ids']) for item in teacher_input_data] + teacher_labels_list = [torch.tensor(item['labels']) for item in teacher_input_data] + + teacher_input_ids_list_2 = [torch.tensor(item['input_ids']) for item in teacher_input_data_2] + teacher_labels_list_2 = [torch.tensor(item['labels']) for item in teacher_input_data_2] + + # Pad sequences to the same length + teacher_input_ids = rnn_utils.pad_sequence(teacher_input_ids_list, batch_first=True) + teacher_labels = rnn_utils.pad_sequence(teacher_labels_list, batch_first=True, padding_value=-100) + + teacher_input_ids_2 = rnn_utils.pad_sequence(teacher_input_ids_list_2, batch_first=True) + teacher_labels_2 = rnn_utils.pad_sequence(teacher_labels_list_2, batch_first=True, padding_value=-100) + + # Create teacher_output dict for both teachers + teacher_output = { + 'teacher_labels': [teacher_labels, teacher_labels_2], + 'teacher_input_ids': [teacher_input_ids, teacher_input_ids_2], + 'teacher_topk_logprobs_group': [topk_data['teacher_topk_logprobs'], topk_data_2['teacher_topk_logprobs']], + 'teacher_topk_indices_group': [topk_data['teacher_topk_indices'], topk_data_2['teacher_topk_indices']], + } + + # 5. Student forward + CTKD backward + # batch is already encoded with student tokenizer + 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 % 10 == 0: + metric = student_model.calculate_metric(is_training=True, adapter_name=ADAPTER_NAME) + logger.info(f'[Step {optim_step}/{MAX_STEPS}] {metric}') + + # 6. 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/cookbook/rl/mopd/gold_npu.py b/cookbook/rl/mopd/gold_npu.py new file mode 100644 index 00000000..303354da --- /dev/null +++ b/cookbook/rl/mopd/gold_npu.py @@ -0,0 +1,304 @@ +import os +from typing import List + +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.model import TransformersModel +from twinkle.loss import GOLDLoss +from twinkle.loss.gold_config import GOLDConfig +from twinkle.sampler import vLLMSampler + +logger = get_logger() + +# ── Configuration ───────────────────────────────────────────────────────────── +STUDENT_MODEL_ID = os.environ.get('STUDENT_MODEL_ID', 'ms://Qwen/Qwen3-0.6B') +TEACHER_MODEL_ID = os.environ.get('TEACHER_MODEL_ID', 'ms://Qwen/qwen2.5-0.5b-instruct') + +MODEL_GPUS = int(os.environ.get('MODEL_GPUS', 1)) +SAMPLER_GPUS = int(os.environ.get('SAMPLER_GPUS', 1)) +NUM_GPUS = MODEL_GPUS + SAMPLER_GPUS + +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-4)) + +GKD_TEMPERATURE = float(os.environ.get('GKD_TEMPERATURE', 1.0)) +ADAPTER_NAME = 'default' + + +# ── NPU Environment Setup ───────────────────────────────────────────────────── + +def setup_npu_env(): + """Setup NPU environment variables for Ascend devices.""" + # 1. 设置 NPU 日志输出 + os.environ["ASCEND_SLOG_PRINT_TO_STDOUT"] = "1" + + # 2. 设置正确的 CANN 路径（必须匹配实际安装路径！） + # 请执行以下命令确认实际路径： + # ls /usr/local/Ascend/ascend-toolkit/ | grep -E '5\.1\.RC[0-9]+' + # 假设输出为 5.1.RC1，则： + os.environ["ASCEND_AICPU_PATH"] = "/usr/local/Ascend/ascend-toolkit/5.1.RC1" + + # 3. 更新 LD_LIBRARY_PATH（包含 lib64 目录） + os.environ["LD_LIBRARY_PATH"] = ( + f"{os.environ.get('LD_LIBRARY_PATH', '')}:" + f"{os.environ['ASCEND_AICPU_PATH']}/lib64" + ) + +def convert_topk_prompt_logprobs( + topk_prompt_logprobs_batch: List, +) -> dict: + """将 vLLM 的 topk_prompt_logprobs 转换为 GOLDLoss 所需的 teacher_output 格式。 + + 这个方法的作用是： + 1. vLLM 返回的 teacher_response 中包含 topk_prompt_logprobs，这是一个嵌套列表结构 + 2. GOLDLoss 需要的是 teacher_topk_logprobs 和 teacher_topk_indices 两个张量 + 3. 本方法将嵌套列表转换为对齐的张量格式，用于知识蒸馏损失计算 + + Args: + topk_prompt_logprobs_batch: 每个输入的 topk_prompt_logprobs 列表。 + 每个元素是 List[Optional[List[(token_id, logprob)]]]， + 形状为 [seq_len, topk]，即每个序列位置有 topk 个 (token_id, logprob) 对。 + + Returns: + 包含以下键的字典： + - 'teacher_topk_logprobs': [batch, seq_len, topk] 形状的张量，存储教师模型的 top-k 对数概率 + - 'teacher_topk_indices': [batch, seq_len, topk] 形状的张量，存储对应的 token ID + """ + batch_logprobs = [] # 存储整个 batch 的对数概率 + batch_indices = [] # 存储整个 batch 的 token 索引 + + # 遍历 batch 中的每个序列 + for seq_topk in topk_prompt_logprobs_batch: + if seq_topk is None: + # 处理 None 情况（没有可用的 prompt logprobs） + # 用占位符填充：单个位置，64 个 top-k 值 + seq_logprobs = [[0.0] * 64] + seq_indices = [[0] * 64] + else: + seq_logprobs = [] + seq_indices = [] + # 遍历序列中的每个位置 + for pos_topk in seq_topk: + if pos_topk is None: + # 第一个位置是 None（因为第一个 token 没有前文），用占位符填充 + seq_logprobs.append([0.0] * 64) + seq_indices.append([0] * 64) + else: + # 提取 (token_id, logprob) 对中的 logprob + seq_logprobs.append([lp for _, lp in pos_topk]) + # 提取 (token_id, logprob) 对中的 token_id + seq_indices.append([tid for tid, _ in pos_topk]) + batch_logprobs.append(seq_logprobs) + batch_indices.append(seq_indices) + + # 将 batch 内的所有序列填充到相同的 seq_len + max_len = max(len(seq) for seq in batch_logprobs) if batch_logprobs else 1 + topk = 64 + + # 对较短的序列进行填充 + for i in range(len(batch_logprobs)): + pad_len = max_len - len(batch_logprobs[i]) + if pad_len > 0: + # 用 0 填充到最大长度 + batch_logprobs[i].extend([[0.0] * topk] * pad_len) + batch_indices[i].extend([[0] * topk] * pad_len) + + # 滚动张量以对齐 labels（第一个位置没有有效的 logprobs，需要移到末尾） + # 这是因为第 i 个位置的 logprobs 是预测第 i+1 个 token 的概率分布 + 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 train(): + """Main training function for MOPD with GOLDLoss on NPU.""" + # Setup NPU environment + setup_npu_env() + + # Initialize device groups + 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, + ) + logger.info(get_device_placement()) + + # ── Student model (trainable) ────────────────────────────────────────────── + student_model = TransformersModel( + model_id=STUDENT_MODEL_ID, + device_mesh=model_mesh, + remote_group='student_model', + ) + + 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=2) + student_model.set_optimizer('AdamW', lr=LEARNING_RATE) + student_model.set_lr_scheduler('CosineWarmupScheduler', num_warmup_steps=5, num_training_steps=MAX_STEPS) + + # ── Configure GOLDLoss ───────────────────────────────────────────────────── + gold_config = GOLDConfig( + # ULD loss configuration + use_uld_loss=True, + use_extended_uld=True, + uld_crossentropy_weight=0.0, # Set > 0 to add cross-entropy loss + uld_distillation_weight=1.0, + uld_student_temperature=GKD_TEMPERATURE, + uld_teacher_temperature=GKD_TEMPERATURE, + uld_skip_student_eos=True, + uld_skip_teacher_eos=True, + # Hybrid loss (optional, for cross-tokenizer distillation) + uld_use_hybrid_loss=True, + beta=0.5, # JSD interpolation coefficient + ) + + # Get tokenizers for ULD alignment + 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) + + loss_fn = GOLDLoss( + config=gold_config, + student_tokenizer=student_tokenizer, + teacher_tokenizer=teacher_tokenizer, + ) + student_model.set_loss(loss_fn, adapter_name=ADAPTER_NAME) + student_model.set_template('Template', model_id=STUDENT_MODEL_ID, adapter_name=ADAPTER_NAME) + + # ── Teacher vLLM sampler (for logits) ────────────────────────────────────── + teacher_sampler = vLLMSampler( + model_id=TEACHER_MODEL_ID, + engine_args={ + 'gpu_memory_utilization': 0.65, + 'max_model_len': 4096, + 'logprobs_mode': 'raw_logprobs', + 'max_logprobs': 64, + }, + device_mesh=sampler_mesh, + remote_group='teacher_sampler', + ) + teacher_sampler.set_template('Template', model_id=TEACHER_MODEL_ID) + + # ── DataLoader (full-text: prompt + response) ────────────────────────────── + dataset = Dataset( + dataset_meta=DatasetMeta( + 'ms://hjh0119/shareAI-Llama3-DPO-zh-en-emoji', + data_slice=range(1000) + ) + ) + dataset.set_template('Template', model_id=STUDENT_MODEL_ID) + dataset.encode() + + dataloader = DataLoader( + dataset=dataset, + batch_size=BATCH_SIZE, + min_batch_size=BATCH_SIZE, + remote_group='student_model', + ) + + logger.info(f'MOPD NPU Training | student={STUDENT_MODEL_ID} teacher={TEACHER_MODEL_ID}') + logger.info(f' T={GKD_TEMPERATURE} batch_size={BATCH_SIZE} lr={LEARNING_RATE}') + + # ── Training Loop ────────────────────────────────────────────────────────── + optim_step = 0 + for batch in dataloader: + if optim_step >= MAX_STEPS: + break + # 在 Ray 模式下，batch 是一个可调用对象，需要调用它来获取实际数据 + if callable(batch): + batch = batch() + + # 1. Teacher computes logits on the full sequences + # max_tokens=0: don't generate new content, just compute logits on input + # Use prompt_logprobs to get teacher logprobs for distillation + + teacher_inputs = [] + for item in batch: + text = student_tokenizer.decode(item['input_ids'], skip_special_tokens=False) + teacher_inputs.append({'messages': [{'role': 'user', 'content': text}]}) + + teacher_response = teacher_sampler.sample( + teacher_inputs, + SamplingParams(max_tokens=0, temperature=1.0, prompt_logprobs=64), + ) + + # 2. Convert teacher response to input format + input_data = [seq.new_input_feature for resp in teacher_response for seq in resp.sequences] + + # 3. Prepare teacher output for GOLDLoss + # GOLDLoss expects teacher_logits, teacher_labels, and teacher_input_ids + # Convert topk_prompt_logprobs to the format expected by GOLDLoss + + # Convert topk_prompt_logprobs to teacher_logits format + topk_data = convert_topk_prompt_logprobs( + [resp.topk_prompt_logprobs for resp in teacher_response] + ) + + # Get teacher input_ids and labels from input_data + # Convert lists to tensors first + # Handle variable sequence lengths by padding to the maximum length + import torch.nn.utils.rnn as rnn_utils + + # Get all input_ids and labels + input_ids_list = [torch.tensor(item['input_ids']) for item in input_data] + labels_list = [torch.tensor(item['labels']) for item in input_data] + + # Pad sequences to the same length + teacher_input_ids = rnn_utils.pad_sequence(input_ids_list, batch_first=True) + teacher_labels = rnn_utils.pad_sequence(labels_list, batch_first=True, padding_value=-100) + + # Create teacher_output dict with topk format + # GOLDLoss has been modified to support topk format + teacher_output = { + 'teacher_labels': teacher_labels, + 'teacher_input_ids': teacher_input_ids, + 'teacher_topk_logprobs': topk_data['teacher_topk_logprobs'], + 'teacher_topk_indices': topk_data['teacher_topk_indices'], + } + + # 4. Student forward + GOLD 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 % 20 == 0: + metric = student_model.calculate_metric(is_training=True, adapter_name=ADAPTER_NAME) + logger.info(f'[Step {optim_step}/{MAX_STEPS}] {metric}') + + # 6. Checkpoint + if optim_step > 0 and optim_step % 100 == 0: + student_model.save(f'mopd-npu-ckpt-{optim_step}', adapter_name=ADAPTER_NAME) + + optim_step += 1 + + # Save final checkpoint + student_model.save('mopd-npu-final', adapter_name=ADAPTER_NAME) + logger.info('MOPD NPU 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 8e1d0e2a..f1f2a5d0 100644 --- a/src/twinkle/loss/__init__.py +++ b/src/twinkle/loss/__init__.py @@ -7,6 +7,8 @@ from .grpo import BNPOLoss, CISPOLoss, DRGRPOLoss, GRPOLoss, GSPOLoss, SAPOLoss from .infonce import InfonceLoss from .mse import MSELoss +from .gold import GOLDLoss +from .ctkd import CTKDLoss torch_loss_mapping = { 'mse': MSELoss, @@ -28,4 +30,6 @@ 'orpo': ORPOLoss, # Embedding / contrastive losses 'infonce': InfonceLoss, + 'gold': GOLDLoss, + 'ctkd': CTKDLoss, } diff --git a/src/twinkle/loss/ctkd.py b/src/twinkle/loss/ctkd.py new file mode 100644 index 00000000..e8628f17 --- /dev/null +++ b/src/twinkle/loss/ctkd.py @@ -0,0 +1,867 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +# https://arxiv.org/pdf/2605.21699 +# X-Token: Projection-Guided Cross-Tokenizer Knowledge Distillation - 2605.21699 +""" +CTKD (Cross-Tokenizer Knowledge Distillation) Loss Implementation. + +This module implements the X-Token approach for knowledge distillation across +models with different tokenizers using a sparse projection matrix. + +Reference: + "X-Token: Projection-Guided Cross-Tokenizer Knowledge Distillation" + (https://arxiv.org/pdf/2605.21699) +""" +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 + +if TYPE_CHECKING: + from transformers import PreTrainedTokenizer + +# Global cache for projection matrices to avoid recomputation +_PROJECTION_MATRIX_CACHE = {} + + +class CTKDLoss(Loss): + """ + 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). + + Example: + >>> from transformers import AutoTokenizer + >>> student_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B") + >>> teacher_tokenizer1 = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf") + >>> teacher_tokenizer2 = AutoTokenizer.from_pretrained("Qwen/Qwen2-7B") + >>> loss_fn = CTKDLoss(student_tokenizer, [teacher_tokenizer1, teacher_tokenizer2]) + >>> # Projection matrices are built automatically during initialization + >>> print(len(loss_fn.projection_matrices)) # 2 (number of teachers) + """ + 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 + 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: + """ + Automatically select loss_type based on vocabulary coverage. + + Coverage is defined as the intersection of student and teacher vocabularies + divided by the union of vocabularies. + + When coverage is high (>= 0.7), use 'hkl' for better efficiency. + When coverage is low (< 0.7), use 'pkl' for better accuracy. + + Returns: + 'pkl' or 'hkl' + """ + # 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 + + + # Check again inside the lock to avoid race condition + if self._projection_matrices_built: + return + + + + + # 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': + 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') + + # Support topk format from vLLM for multiple teachers + 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 + + # Create full logits tensor initialized with very negative values + teacher_logits = torch.full( + (batch_size, seq_len, vocab_size), + -1e10, # Very negative value for log-space + dtype=teacher_topk_logprobs.dtype, + device=student_logits.device if student_logits is not None else teacher_topk_logprobs.device + ) + + # Scatter the topk logprobs into the full logits tensor + teacher_logits.scatter_( + dim=2, + index=teacher_topk_indices.to(teacher_logits.device), + src=teacher_topk_logprobs.to(teacher_logits.device) + ) + teacher_logits_group.append(teacher_logits) + + 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 + 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) + + num_tokens = labels.ne(-100).sum().item() if labels is not None else 0 + return LossOutput(loss=total_loss, num_tokens=num_tokens) + + + 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 + student_indices_tensor = torch.tensor(student_indices, dtype=torch.long) + teacher_indices_tensor = torch.tensor(teacher_indices, dtype=torch.long) + values_tensor = torch.tensor(values, dtype=torch.float32).half() # Use half precision + + # 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() * 2) / (1024**2) + + # Move to specified device if provided + if self.device is not None: + self.projection_student_indices_list[teacher_index] = self.projection_student_indices_list[teacher_index].to(self.device) + self.projection_teacher_indices_list[teacher_index] = self.projection_teacher_indices_list[teacher_index].to(self.device) + self.projection_values_list[teacher_index] = self.projection_values_list[teacher_index].to(self.device) + + 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_logits: torch.Tensor, + 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) + + This preserves the full distribution information and is suitable when + critical tokens (e.g., numbers in math tasks) don't have exact matches. + + Args: + student_logits: [batch, seq_len, student_vocab_size] + teacher_logits: [batch, seq_len, teacher_vocab_size] + labels: [batch, seq_len] + + Returns: + Scalar loss value. + """ + # Shift logits and labels for next-token prediction + shift_student_logits = student_logits[..., :-1, :].contiguous() + shift_teacher_logits = teacher_logits[..., :-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 probabilities with temperature + student_probs = F.softmax(shift_student_logits / self.temperature, dim=-1) + teacher_probs = F.softmax(shift_teacher_logits / self.temperature, dim=-1) + + # Align sequence lengths by taking the minimum + student_seq_len = student_probs.shape[1] + teacher_seq_len = 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, :] + teacher_probs = 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 + + batch_size, seq_len, student_vocab = student_probs.shape + # Use actual teacher vocabulary size from teacher_probs tensor + _, _, teacher_vocab_size = 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) + ) + + # Normalize projected probabilities + projected_student_probs = projected_student_probs / ( + projected_student_probs.sum(dim=-1, keepdim=True) + 1e-8 + ) + + # Debug: Check dimensions before KL divergence + if projected_student_probs.size(-1) != 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 {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 {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." + ) + + # Compute KL divergence: KL(p_T || p̃_S) + # KL(P||Q) = Σ P(x) * log(P(x)/Q(x)) + # Using F.kl_div: input=log(Q), target=P, reduction='none' + log_projected_student = torch.log(projected_student_probs + 1e-8) + kl_div = F.kl_div( + log_projected_student, + teacher_probs, + reduction='none' + ).sum(dim=-1) # Sum over vocabulary dimension + + # Apply mask and average + masked_kl = kl_div * loss_mask + loss = masked_kl.sum() / (loss_mask.sum() + 1e-8) + + return loss + + + def _compute_hkl_loss( + self, + student_logits: torch.Tensor, + teacher_logits: torch.Tensor, + labels: torch.Tensor, + teacher_index: int = 0, + ) -> torch.Tensor: + """ + Compute H-KL (Heuristic KL) loss. + + H-KL uses the best teacher token mapping for each student token: + t* = argmax_{t'} W[s, t'] + + This creates a one-to-one mapping and is suitable when most critical + tokens have exact matches in both vocabularies. + + Args: + student_logits: [batch, seq_len, student_vocab_size] + teacher_logits: [batch, seq_len, teacher_vocab_size] + labels: [batch, seq_len] + + Returns: + Scalar loss value. + """ + # Shift logits and labels for next-token prediction + shift_student_logits = student_logits[..., :-1, :].contiguous() + shift_teacher_logits = teacher_logits[..., :-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 probabilities with temperature + student_probs = F.softmax(shift_student_logits / self.temperature, dim=-1) + teacher_probs = F.softmax(shift_teacher_logits / self.temperature, dim=-1) + + # Align sequence lengths by taking the minimum + student_seq_len = student_probs.shape[1] + teacher_seq_len = 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, :] + teacher_probs = 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 = 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) + ) + + # Normalize + mapped_student_probs = mapped_student_probs / ( + mapped_student_probs.sum(dim=-1, keepdim=True) + 1e-8 + ) + + # Compute KL divergence: KL(p_T || mapped_p_S) + log_mapped_student = torch.log(mapped_student_probs + 1e-8) + kl_div = F.kl_div( + log_mapped_student, + teacher_probs, + reduction='none' + ).sum(dim=-1) + + # Apply mask and average + masked_kl = kl_div * loss_mask + loss = masked_kl.sum() / (loss_mask.sum() + 1e-8) + + 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 diff --git a/src/twinkle/loss/gold.py b/src/twinkle/loss/gold.py new file mode 100644 index 00000000..4acf4f2d --- /dev/null +++ b/src/twinkle/loss/gold.py @@ -0,0 +1,544 @@ +from typing import List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from twinkle.data_format import LossOutput +from twinkle.loss.base import Loss +from twinkle.loss.gold_config import GOLDConfig + + +class GOLDLoss(Loss): + + def __init__(self, config: GOLDConfig, student_tokenizer=None, teacher_tokenizer=None, + device: Optional[torch.device] = None, ): + self.device = device + self.crossentropy_weight = config.uld_crossentropy_weight + self.distillation_weight = config.uld_distillation_weight + self.student_temperature = config.uld_student_temperature + self.teacher_temperature = config.uld_teacher_temperature + self.skip_student_eos = config.uld_skip_student_eos + self.skip_teacher_eos = config.uld_skip_teacher_eos + self.use_extended_uld = config.use_extended_uld + self.ignore_index = -100 + + # Add tokenizers for enhanced alignment + self.student_tokenizer = student_tokenizer + self.teacher_tokenizer = teacher_tokenizer + + # Hybrid ULD configuration + self.use_hybrid_loss = getattr(config, "uld_use_hybrid_loss", False) + self.hybrid_matched_weight = getattr(config, "uld_hybrid_matched_weight", None) + self.hybrid_unmatched_weight = getattr(config, "uld_hybrid_unmatched_weight", None) + self.beta = getattr(config, "beta", 1.0) # For JSD loss in hybrid matched tokens + + # Initialize vocabulary mapping for hybrid loss + self._vocab_mapping = None + self._teacher_matched_ids = None + self._student_matched_ids = None + if self.use_hybrid_loss and student_tokenizer is not None and teacher_tokenizer is not None: + self._initialize_vocabulary_mapping() + + def __call__( + self, + inputs, + outputs, + **kwargs + ) -> LossOutput: + # Extract required tensors from inputs and outputs + # inputs should contain student_input_ids, teacher_input_ids + # outputs should contain student_logits, teacher_logits + # labels should be in inputs or outputs + + # 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 and labels from kwargs or inputs/outputs + teacher_logits = kwargs.get('teacher_logits') + if teacher_logits is None: + teacher_logits = outputs.get('teacher_logits') + + # Support topk format from vLLM + teacher_topk_logprobs = kwargs.get('teacher_topk_logprobs') + teacher_topk_indices = kwargs.get('teacher_topk_indices') + + # If we have topk format but not full logits, convert topk to logits + if teacher_logits is None and teacher_topk_logprobs is not None and teacher_topk_indices is not None: + # Get vocabulary size from teacher tokenizer if available, otherwise fallback to student + vocab_size = len(self.teacher_tokenizer) if self.teacher_tokenizer is not None else student_logits.size(-1) + batch_size, seq_len, topk = teacher_topk_logprobs.shape + # Create full logits tensor initialized with very negative values + teacher_logits = torch.full( + (batch_size, seq_len, vocab_size), + -1e10, # Very negative value for log-space + dtype=teacher_topk_logprobs.dtype, + device=teacher_topk_logprobs.device + ) + + # Scatter the topk logprobs into the full logits tensor + teacher_logits.scatter_( + dim=2, + index=teacher_topk_indices, + src=teacher_topk_logprobs + ) + + if teacher_logits is None: + raise ValueError("Teacher logits not found in kwargs or outputs. " + "Provide either teacher_logits or (teacher_topk_logprobs + teacher_topk_indices)") + + teacher_labels = kwargs.get('teacher_labels') + if teacher_labels is None: + teacher_labels = inputs.get('teacher_labels') + if teacher_labels is None: + raise ValueError("Teacher labels not found in kwargs or inputs") + + student_input_ids = inputs.get('input_ids') + if student_input_ids is None: + raise ValueError("Student input_ids not found in inputs") + + teacher_input_ids = kwargs.get('teacher_input_ids') + if teacher_input_ids is None: + teacher_input_ids = inputs.get('teacher_input_ids') + if teacher_input_ids is None: + raise ValueError("Teacher input_ids not found in kwargs or inputs") + + # Compute cross-entropy loss for student + if self.crossentropy_weight > 0: + shift_logits = student_logits[..., :-1, :].contiguous() + shift_labels = student_labels[..., 1:].contiguous() + loss_fct = nn.CrossEntropyLoss(ignore_index=self.ignore_index) + crossentropy_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) + crossentropy_loss = self.crossentropy_weight * crossentropy_loss + else: + crossentropy_loss = 0.0 + + # Compute distillation loss using ULD approximation + distillation_loss = self._compute_distillation_loss( + student_logits, teacher_logits, student_labels, teacher_labels, student_input_ids, teacher_input_ids + ) + + total_loss = crossentropy_loss + distillation_loss + + return LossOutput(loss=total_loss, num_tokens=student_labels.ne(self.ignore_index).sum()) + + def _initialize_vocabulary_mapping(self): + """Initialize vocabulary mapping for hybrid ULD loss.""" + # Computing vocabulary mapping for hybrid ULD + + student_vocab = self.student_tokenizer.get_vocab() + teacher_vocab = self.teacher_tokenizer.get_vocab() + + # Create reverse mapping for student + student_token_to_id = dict(student_vocab.items()) + + vocab_mapping = {} + teacher_matched_ids = set() + student_matched_ids = set() + + for token_str, teacher_id in teacher_vocab.items(): + if token_str in student_token_to_id: + student_id = student_token_to_id[token_str] + vocab_mapping[teacher_id] = student_id + teacher_matched_ids.add(teacher_id) + student_matched_ids.add(student_id) + + self._vocab_mapping = vocab_mapping + self._teacher_matched_ids = teacher_matched_ids + self._student_matched_ids = student_matched_ids + + max_matched_teacher_id = max(self._vocab_mapping.keys()) + self.mapping_tensor = torch.full((max_matched_teacher_id + 1,), -1, dtype=torch.long) # -1 for unmapped ids + for k, v in self._vocab_mapping.items(): + self.mapping_tensor[k] = v + if self.device is not None: + self.mapping_tensor = self.mapping_tensor.to(self.device) + + def _compute_cross_entropy( + self, + student_logits: torch.Tensor, + student_labels: torch.Tensor + ) -> torch.Tensor: + """计算cross-entropy loss""" + if self.crossentropy_weight <= 0: + return torch.tensor(0.0, device=student_logits.device, requires_grad=True) + + shift_logits = student_logits[..., :-1, :].contiguous() + shift_labels = student_labels[..., 1:].contiguous() + + loss_fct = nn.CrossEntropyLoss(ignore_index=self.ignore_index) + ce_loss = loss_fct( + shift_logits.view(-1, shift_logits.size(-1)), + shift_labels.view(-1) + ) + return self.crossentropy_weight * ce_loss + + def _compute_distillation_loss( + self, + student_logits: torch.Tensor, + teacher_logits: torch.Tensor, + student_labels: torch.Tensor, + teacher_labels: torch.Tensor, + student_input_ids: torch.Tensor, + teacher_input_ids: torch.Tensor, + ) -> torch.Tensor: + """计算ULD蒸馏损失""" + # 获取答案区域 + student_answer_idx, student_answer_size = self._get_answer_regions(student_labels) + teacher_answer_idx, teacher_answer_size = self._get_answer_regions(teacher_labels) + + if self.skip_student_eos: + student_answer_size = [s - 1 for s in student_answer_size] + if self.skip_teacher_eos: + teacher_answer_size = [t - 1 for t in teacher_answer_size] + + # 边界检查 + if not student_answer_size or not teacher_answer_size: + return torch.zeros(1, device=student_logits.device, requires_grad=True) * 1e-8 + + batch_size = student_logits.size(0) + distillation_losses = [] + + for i in range(batch_size): + s_start = student_answer_idx[i] + s_size = student_answer_size[i] + t_start = teacher_answer_idx[i] + t_size = teacher_answer_size[i] + + if s_size <= 0 or t_size <= 0: + loss_i = student_logits[i].sum() * 0.0 + # Ensure the loss tensor requires gradients + loss_i = loss_i.detach().requires_grad_(True) + distillation_losses.append(loss_i) + continue + + # 提取答案logits + student_ans_logits = student_logits[i, s_start:s_start + s_size] + teacher_ans_logits = teacher_logits[i, t_start:t_start + t_size] + + # 转换为概率 + student_probs = F.softmax(student_ans_logits / self.student_temperature, dim=-1) + teacher_probs = F.softmax(teacher_ans_logits / self.teacher_temperature, dim=-1) + + # 获取token IDs用于对齐 + + student_token_ids = student_input_ids[i, s_start:s_start + s_size].tolist() + teacher_token_ids = teacher_input_ids[i, t_start:t_start + t_size].tolist() + + # Token对齐 + if self.use_extended_uld: + student_groups, teacher_groups = self._build_alignment_groups_from_ids( + student_token_ids, teacher_token_ids + ) + + student_aligned = self._merge_probabilities_with_groups( + student_probs, student_groups, student_token_ids + ) + teacher_aligned = self._merge_probabilities_with_groups( + teacher_probs, teacher_groups, teacher_token_ids + ) + else: + min_len = min(len(student_token_ids), len(teacher_token_ids)) + student_aligned = student_probs[:min_len] + teacher_aligned = teacher_probs[:min_len] + + # 计算损失 + if self.use_hybrid_loss and self._vocab_mapping: + aligned_loss = self._compute_hybrid_uld_loss(student_aligned, teacher_aligned) + else: + aligned_loss = self._compute_basic_uld_loss(student_aligned, teacher_aligned) + + distillation_losses.append(aligned_loss) + distillation_loss = torch.stack(distillation_losses).mean() + return self.distillation_weight * distillation_loss + + def _get_answer_regions(self, labels: torch.Tensor) -> Tuple[List[int], List[int]]: + """获取答案区域的起始位置和大小""" + indices = [] + sizes = [] + + for label in labels: + mask = label.ne(self.ignore_index) + if not mask.any(): + indices.append(0) + sizes.append(0) + continue + + valid_indices = mask.nonzero(as_tuple=True)[0] + indices.append(int(valid_indices[0].item())) + sizes.append(int(mask.sum().item())) + + return indices, sizes + + def _build_alignment_groups_from_ids( + self, + student_token_ids: List[int], + teacher_token_ids: List[int] + ) -> Tuple[List[List[int]], List[List[int]]]: + """ + 基于文本内容构建对齐组 + 使用贪心子串匹配算法 + """ + + def decode_tokens(tokenizer, token_ids): + pieces = [] + prev = "" + for k in range(len(token_ids)): + cur = tokenizer.decode(token_ids[:k + 1], skip_special_tokens=False) + pieces.append(cur[len(prev):]) + prev = cur + return pieces + + student_pieces = decode_tokens(self.student_tokenizer, student_token_ids) + teacher_pieces = decode_tokens(self.teacher_tokenizer, teacher_token_ids) + + # 贪心匹配算法 + student_groups = [] + teacher_groups = [] + s_idx = 0 + t_idx = 0 + + while s_idx < len(student_pieces) and t_idx < len(teacher_pieces): + student_text = "" + teacher_text = "" + student_group = [] + teacher_group = [] + + # 尝试找到最短的连续匹配序列 + while s_idx < len(student_pieces) and t_idx < len(teacher_pieces): + if not student_group: + student_group.append(s_idx) + student_text += student_pieces[s_idx] + s_idx += 1 + + if not teacher_group: + teacher_group.append(t_idx) + teacher_text += teacher_pieces[t_idx] + t_idx += 1 + + # 检查是否匹配 + if student_text == teacher_text: + student_groups.append(student_group) + teacher_groups.append(teacher_group) + break + elif len(student_text) < len(teacher_text): + if s_idx < len(student_pieces): + student_group.append(s_idx) + student_text += student_pieces[s_idx] + s_idx += 1 + else: + break + else: + if t_idx < len(teacher_pieces): + teacher_group.append(t_idx) + teacher_text += teacher_pieces[t_idx] + t_idx += 1 + else: + break + else: + # 未完全匹配，添加剩余部分 + if student_group and teacher_group: + student_groups.append(student_group) + teacher_groups.append(teacher_group) + + return student_groups, teacher_groups + + def _merge_probabilities_with_groups( + self, + probs: torch.Tensor, + alignment_groups: List[List[int]], + token_ids: List[int], + ) -> torch.Tensor: + """ + 根据对齐组合并概率分布 + 使用链式法则: P_merged = P(y|x_0) * P(x_1|x_0) * P(x_2|x_0,x_1) * ... + """ + aligned_probs = [] + + for group in alignment_groups: + if len(group) > 1: + # 第一个token的边际概率 + marginal_probs = probs[group[0]] # [vocab_size] + + # 后续token的条件概率（标量） + conditional_product = 1.0 + for k in range(1, len(group)): + cond_prob = probs[group[k], token_ids[group[k - 1]]] + conditional_product *= cond_prob + + merged_probs = marginal_probs * conditional_product + aligned_probs.append(merged_probs) + elif len(group) == 1: + aligned_probs.append(probs[group[0]]) + + if aligned_probs: + return torch.stack(aligned_probs) + else: + # 返回一个空的但需要梯度的张量 + empty_tensor = probs[:0].detach().requires_grad_(True) + return empty_tensor + + def _compute_basic_uld_loss( + self, + student_aligned: torch.Tensor, + teacher_aligned: torch.Tensor, + ) -> torch.Tensor: + """基础ULD损失：排序后的L1距离""" + student_sorted = student_aligned.sort(dim=-1, descending=True).values + teacher_sorted = teacher_aligned.sort(dim=-1, descending=True).values + + # Padding到相同vocab size + s_vocab = student_sorted.size(-1) + t_vocab = teacher_sorted.size(-1) + max_vocab = max(s_vocab, t_vocab) + + if s_vocab < max_vocab: + student_sorted = F.pad(student_sorted, (0, max_vocab - s_vocab)) + if t_vocab < max_vocab: + teacher_sorted = F.pad(teacher_sorted, (0, max_vocab - t_vocab)) + + teacher_sorted = teacher_sorted.to(student_sorted.device) + + loss = F.l1_loss(student_sorted, teacher_sorted, reduction="sum") + loss /= student_aligned.size(0) + + return loss + + def _compute_hybrid_uld_loss( + self, + student_aligned: torch.Tensor, + teacher_aligned: torch.Tensor, + ) -> torch.Tensor: + """混合ULD损失：matched用JSD，unmatched用排序L1""" + # 统一设备：确保所有tensor都在同一个设备上 + device = student_aligned.device + # 确保teacher_aligned也在同一设备上 + teacher_aligned = teacher_aligned.to(device) + # 确保mapping_tensor在同一设备上 + mapping_tensor = self.mapping_tensor.to(device) + + s_vocab = student_aligned.size(-1) + t_vocab = teacher_aligned.size(-1) + + # 创建matched/unmatched masks + if self._teacher_matched_ids: + teacher_matched_idx = torch.tensor( + sorted(self._teacher_matched_ids), dtype=torch.long, device=device + ) + student_matched_idx = mapping_tensor[teacher_matched_idx] + else: + teacher_matched_idx = torch.tensor([], dtype=torch.long, device=device) + student_matched_idx = torch.tensor([], dtype=torch.long, device=device) + + teacher_matched_mask = torch.zeros(t_vocab, dtype=torch.bool, device=device) + student_matched_mask = torch.zeros(s_vocab, dtype=torch.bool, device=device) + + if len(teacher_matched_idx) > 0: + teacher_matched_mask[teacher_matched_idx] = True + student_matched_mask[student_matched_idx] = True + + # 1. Matched tokens的JSD损失 + matched_loss = torch.tensor(0.0, device=device, requires_grad=True) + matched_count = 0 + + if len(teacher_matched_idx) > 0: + teacher_matched_probs = teacher_aligned[:, teacher_matched_idx] + student_matched_probs = student_aligned[:, student_matched_idx] + matched_count = teacher_matched_probs.size(-1) + + matched_loss = self._compute_jsd_for_matched( + student_matched_probs, teacher_matched_probs + ) + # 2. Unmatched tokens的排序L1损失 + teacher_unmatched = teacher_aligned[:, ~teacher_matched_mask] + student_unmatched = student_aligned[:, ~student_matched_mask] + + unmatched_loss = torch.tensor(0.0, device=device, requires_grad=True) + if teacher_unmatched.size(-1) > 0 and student_unmatched.size(-1) > 0: + teacher_sorted = teacher_unmatched.sort(dim=-1, descending=True).values + student_sorted = student_unmatched.sort(dim=-1, descending=True).values + + t_size = teacher_sorted.size(-1) + s_size = student_sorted.size(-1) + max_size = max(t_size, s_size) + + if t_size < max_size: + teacher_sorted = F.pad(teacher_sorted, (0, max_size - t_size)) + if s_size < max_size: + student_sorted = F.pad(student_sorted, (0, max_size - s_size)) + + unmatched_loss = F.l1_loss(student_sorted, teacher_sorted, reduction="sum") + unmatched_loss /= student_aligned.size(0) + + # 3. 加权组合 + if self.hybrid_matched_weight is None: + w_matched = matched_count / max(1, t_vocab) + w_unmatched = 1.0 - w_matched + else: + w_matched = self.hybrid_matched_weight + w_unmatched = self.hybrid_unmatched_weight + + total_loss = w_matched * matched_loss + w_unmatched * unmatched_loss + + # 保存用于logging + self.last_matched_loss = matched_loss + self.last_unmatched_loss = unmatched_loss + + return total_loss + + def _compute_jsd_for_matched( + self, + student_probs: torch.Tensor, + teacher_probs: torch.Tensor, + epsilon: float = 1e-8 + ) -> torch.Tensor: + """计算matched tokens的JSD损失，添加数值稳定性处理""" + # 统一设备：确保所有tensor都在同一个设备上 + device = student_probs.device + teacher_probs = teacher_probs.to(device) + + batch_seq_len, num_matched = student_probs.shape + + # 检查输入概率分布是否有效 + if torch.isnan(student_probs).any() or torch.isnan(teacher_probs).any(): + return torch.tensor(0.0, device=device, requires_grad=True) + + # 添加epsilon防止数值下溢和log(0) + student_probs = student_probs.clamp(min=epsilon) + teacher_probs = teacher_probs.clamp(min=epsilon) + + # 重新归一化概率分布 + student_probs = student_probs / student_probs.sum(dim=-1, keepdim=True) + teacher_probs = teacher_probs / teacher_probs.sum(dim=-1, keepdim=True) + + student_flat = student_probs.view(-1, num_matched) + teacher_flat = teacher_probs.view(-1, num_matched) + + # JSD = 0.5 * KL(P||M) + 0.5 * KL(Q||M), where M = 0.5*(P+Q) + m = 0.5 * (student_flat + teacher_flat) + + # 添加epsilon到中间分布 + m = m.clamp(min=epsilon) + m = m / m.sum(dim=-1, keepdim=True) + + # 直接对概率分布取对数，添加epsilon防止数值问题 + log_m = torch.log(m + epsilon) + log_student = torch.log(student_flat + epsilon) + log_teacher = torch.log(teacher_flat + epsilon) + + # 使用log_target=True，传入log概率 + kl_p_m = F.kl_div(log_m, log_student, reduction='batchmean', log_target=True) + kl_q_m = F.kl_div(log_m, log_teacher, reduction='batchmean', log_target=True) + jsd = 0.5 * (kl_p_m + kl_q_m) + + # 检查结果是否有效 + if torch.isnan(jsd) or torch.isinf(jsd): + return torch.tensor(0.0, device=student_probs.device, requires_grad=True) + + return jsd \ No newline at end of file diff --git a/src/twinkle/loss/gold_config.py b/src/twinkle/loss/gold_config.py new file mode 100644 index 00000000..63313215 --- /dev/null +++ b/src/twinkle/loss/gold_config.py @@ -0,0 +1,393 @@ +from dataclasses import dataclass, field +from typing import Any + + +@dataclass +class GOLDConfig(): + r""" + Configuration class for [`GOLDTrainer`]. + + This class includes only the parameters that are specific to GOLD training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] and [`SFTConfig`] documentation. + + Args: + temperature (`float`, *optional*, defaults to `0.9`): + Temperature for sampling. The higher the temperature, the more random the completions. + lmbda (`float`, *optional*, defaults to `0.5`): + Lambda parameter that controls the student data fraction (i.e., the proportion of on-policy + student-generated outputs). + beta (`float`, *optional*, defaults to `0.5`): + Interpolation coefficient between `0.0` and `1.0` of the Generalized Jensen-Shannon Divergence loss. When + beta is `0.0`, the loss is the KL divergence. When beta is `1.0`, the loss is the Inverse KL Divergence. + max_completion_length (`int`, *optional*, defaults to `128`): + Maximum number of tokens to generate per completion. + teacher_model_name_or_path (`str`, *optional*): + Model name or path of the teacher model. If `None`, the teacher model will be the same as the model being + trained. + teacher_model_revision (`str` or `None`, *optional*, defaults to `None`): + Model revision of the teacher model (e.g., branch name, tag, or commit hash). If `None`, the default + revision is used. + teacher_model_init_kwargs (`dict[str, Any]`, *optional*): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the teacher model + from a string. + teacher_tokenizer_name_or_path (`str`, *optional*): + Tokenizer name or path for the teacher model. If None when using ULD loss, will use the same tokenizer as + the student model (not recommended for cross-tokenizer distillation). + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model. + seq_kd (`bool`, *optional*, defaults to `False`): + Seq_kd parameter that controls whether to perform Sequence-Level KD (can be viewed as supervised FT on + teacher-generated output). + num_generations (`int`, *optional*, defaults to `1`): + Number of generations per prompt. Each prompt is repeated this many times in the generation batch. + generation_batch_size (`int` or `None`, *optional*, defaults to `None`): + Number of unique prompts per worker per optimizer step. If `None`, it is computed from + `(per_device_train_batch_size * gradient_accumulation_steps) // num_generations`. + use_uld_loss (`bool`, *optional*, defaults to `False`): + Whether to use Universal Logit Distillation (ULD) loss instead of Generalized Jensen-Shannon Divergence + loss. + uld_crossentropy_weight (`float`, *optional*, defaults to `0.0`): + Weight for the cross-entropy loss component in ULD loss. If 0, only ULD distillation loss is used. + uld_distillation_weight (`float`, *optional*, defaults to `1.0`): + Weight for the distillation loss component in ULD loss. + uld_student_temperature (`float`, *optional*, defaults to `1.0`): + Temperature for student logits in ULD loss computation. + uld_teacher_temperature (`float`, *optional*, defaults to `1.0`): + Temperature for teacher logits in ULD loss computation. + uld_skip_student_eos (`bool`, *optional*, defaults to `True`): + Whether to skip EOS token for student in ULD loss computation. + uld_skip_teacher_eos (`bool`, *optional*, defaults to `True`): + Whether to skip EOS token for teacher in ULD loss computation. + use_vllm (`bool`, *optional*, defaults to `False`): + Whether to use vLLM for generating completions from the student model. Requires `vllm` to be installed. + vllm_mode (`str`, *optional*, defaults to `"colocate"`): + Mode for student vLLM integration. Either `"server"` (connect to a running TRL vLLM server) or `"colocate"` + (run vLLM in the same process). + vllm_server_host (`str`, *optional*, defaults to `"0.0.0.0"`): + Host of the vLLM server for the student model (if `vllm_mode="server"`). + vllm_server_port (`int`, *optional*, defaults to `8001`): + Port of the vLLM server for the student model (if `vllm_mode="server"`). + vllm_server_timeout (`float`, *optional*, defaults to `240.0`): + Timeout for connecting to the student vLLM server (if `vllm_mode="server"`). + vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.9`): + GPU memory utilization for the colocated student vLLM engine (if `vllm_mode="colocate"`). It is recommended + to set this to a low value if the student and teacher models share the same GPU. + vllm_tensor_parallel_size (`int`, *optional*, defaults to `1`): + Tensor parallel size for the colocated student vLLM engine (if `vllm_mode="colocate"`). + vllm_structured_outputs_regex (`str`, *optional*): + Regex for vLLM structured outputs for the student model. + vllm_server_base_url (`str`, *optional*): + Base URL for the vLLM server (e.g., `"http://localhost:8001"`). If provided, `vllm_server_host` and + `vllm_server_port` are ignored. + vllm_group_port (`int`, *optional*, defaults to `51216`): + Port for the vLLM weight-update group (NCCL communicator). Unless the port is occupied, there is no need to + change it. + vllm_max_model_length (`int`, *optional*): + Maximum model sequence length for the colocated vLLM engine when `vllm_mode="colocate"`. Defaults to the + model's maximum context length. + vllm_model_impl (`str`, *optional*, defaults to `"vllm"`): + Model implementation backend to use in vLLM. Use `"vllm"` (default) or `"transformers"`. + vllm_sync_frequency (`int`, *optional*, defaults to `1`): + Frequency (in training steps) to synchronize student model weights to vLLM engine. Set to 1 to sync after + every step. + vllm_enable_sleep_mode (`bool`, *optional*, defaults to `False`): + Enable vLLM sleep mode to offload student weights/cache during the optimizer step. Keeps GPU memory usage + low, but waking the engine adds host–device transfer latency. + """ + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=1e-7, + metadata={"help": "The initial learning rate for AdamW."}, + ) + + # GOLD-specific parameters + temperature: float = field( + default=0.9, + metadata={"help": "Temperature for sampling. The higher the temperature, the more random the completions."}, + ) + top_p: float = field( + default=0.95, + metadata={ + "help": "If set to float < 1, only the smallest set of most probable tokens with probabilities that add up to " + "`top_p` or higher are kept for generation." + }, + ) + top_k: int = field( + default=0, + metadata={ + "help": "Number of highest probability vocabulary tokens to keep for top-k-filtering. If `0`, " + "top-k-filtering is disabled and all tokens are considered." + }, + ) + lmbda: float = field( + default=0.5, + metadata={ + "help": "Lambda parameter that controls the student data fraction (i.e., the proportion of on-policy " + "student-generated outputs)." + }, + ) + beta: float = field( + default=0.5, + metadata={ + "help": "Interpolation coefficient between `0.0` and `1.0` of the Generalized Jensen-Shannon Divergence " + "loss. When beta is `0.0`, the loss is the KL divergence. When beta is `1.0`, the loss is the Inverse KL " + "Divergence." + }, + ) + max_completion_length: int = field( + default=128, + metadata={"help": "Maximum number of tokens to generate per completion."}, + ) + teacher_model_name_or_path: str | None = field( + default=None, + metadata={ + "help": "Model name or path of the teacher model. If `None`, the teacher model will be the same as the " + "model being trained." + }, + ) + teacher_model_revision: str | None = field( + default=None, + metadata={ + "help": "Model revision of the teacher model (e.g., branch name, tag, or commit hash). If `None`, the " + "default revision is used." + }, + ) + teacher_model_init_kwargs: dict[str, Any] | str | None = field( + default=None, + metadata={ + "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the " + "teacher model from a string." + }, + ) + teacher_tokenizer_name_or_path: str | None = field( + default=None, + metadata={ + "help": "Tokenizer name or path for the teacher model. If None when using ULD loss, will use the same " + "tokenizer as the student model (not recommended for cross-tokenizer distillation)." + }, + ) + disable_dropout: bool = field( + default=True, + metadata={"help": "Whether to disable dropouts in `model`."}, + ) + seq_kd: bool = field( + default=False, + metadata={ + "help": "Seq_kd parameter that controls whether to perform Sequence-Level KD (can be viewed as supervised " + "FT on teacher-generated output)." + }, + ) + num_generations: int = field( + default=1, + metadata={ + "help": "Number of generations per prompt. Increasing this will decrease the number of unique prompts per optimization step." + }, + ) + generation_batch_size: int | None = field( + default=None, + metadata={ + "help": "Number of unique prompts per worker per optimizer step. " + "If None, computed from (per_device_train_batch_size * gradient_accumulation_steps) // num_generations." + }, + ) + + # ULD Loss parameters + use_uld_loss: bool = field( + default=False, + metadata={ + "help": "Whether to use Universal Logit Distillation (ULD) loss instead of Generalized Jensen-Shannon Divergence loss." + }, + ) + use_extended_uld: bool = field( + default=True, + metadata={ + "help": ( + "Whether to enable extended ULD alignment that uses tokenizers to align and merge token " + "probabilities across student and teacher tokenizations. When True, the trainer will compute " + "token mappings and merge probabilities for split tokens; when False, ULD will use simple " + "positional truncation like in the original ULD paper." + ) + }, + ) + uld_use_hybrid_loss: bool = field( + default=False, + metadata={ + "help": ( + "Whether to use a hybrid loss that combines ULD loss and JSD loss. When True, the final loss is a " + "a combination of JSD for known token mappings and ULD for unknown token mappings." + ) + }, + ) + uld_hybrid_matched_weight: float | None = field( + default=None, + metadata={ + "help": ( + "Weight for the matched token loss component when using hybrid ULD + JSD loss. This weight scales " + "the JSD loss computed over tokens that have a direct mapping between student and teacher " + "tokenizations. If None, uses adaptive weighting based on vocabulary overlap. Must be set together " + "with uld_hybrid_unmatched_weight (both None or both float)." + + ) + }, + ) + uld_hybrid_unmatched_weight: float | None = field( + default=None, + metadata={ + "help": ( + "Weight for the unmatched token loss component when using hybrid ULD + JSD loss. This weight scales " + "the ULD loss computed over tokens that do not have a direct mapping between student and teacher " + "tokenizations. If None, uses adaptive weighting based on vocabulary overlap. Must be set together " + "with uld_hybrid_matched_weight (both None or both float)." + ) + }, + ) + uld_crossentropy_weight: float = field( + default=0.0, + metadata={"help": "Weight for the cross-entropy loss component in ULD loss."}, + ) + uld_distillation_weight: float = field( + default=1.0, + metadata={"help": "Weight for the distillation loss component in ULD loss."}, + ) + uld_student_temperature: float = field( + default=1.0, + metadata={"help": "Temperature for student logits in ULD loss computation."}, + ) + uld_teacher_temperature: float = field( + default=1.0, + metadata={"help": "Temperature for teacher logits in ULD loss computation."}, + ) + + uld_skip_student_eos: bool = field( + default=True, + metadata={"help": "Whether to skip EOS token for student in ULD loss computation."}, + ) + uld_skip_teacher_eos: bool = field( + default=True, + metadata={"help": "Whether to skip EOS token for teacher in ULD loss computation."}, + ) + + # vLLM parameters + use_vllm: bool = field( + default=False, + metadata={"help": "Whether to use vLLM for generating completions. Requires `vllm` to be installed."}, + ) + vllm_mode: str = field( + default="colocate", + metadata={ + "help": 'Mode for vLLM integration. Either "server" (connect to a running TRL vLLM server) or "colocate" (run vLLM in the same process).' + }, + ) + vllm_server_base_url: str | None = field( + default=None, + metadata={ + "help": 'Base URL for the vLLM server (e.g., "http://localhost:8001"). If provided, vllm_server_host and vllm_server_port are ignored.' + }, + ) + vllm_server_host: str = field( + default="0.0.0.0", + metadata={"help": 'Host of the vLLM server when `vllm_mode="server"`.'}, + ) + vllm_server_port: int = field( + default=8001, + metadata={"help": 'Port of the vLLM server when `vllm_mode="server"`.'}, + ) + vllm_server_timeout: float = field( + default=240.0, + metadata={"help": 'Timeout (in seconds) for connecting to the vLLM server when `vllm_mode="server"`.'}, + ) + vllm_group_port: int = field( + default=51216, + metadata={"help": "Port for the vLLM weight-update group (NCCL communicator)."}, + ) + vllm_gpu_memory_utilization: float = field( + default=0.9, + metadata={ + "help": 'GPU memory utilization for the colocated vLLM engine when `vllm_mode="colocate"`. Lower values reduce contention when sharing a device with the student/teacher models.' + }, + ) + vllm_tensor_parallel_size: int = field( + default=1, + metadata={"help": 'Tensor parallel size for the colocated vLLM engine when `vllm_mode="colocate"`.'}, + ) + vllm_max_model_length: int | None = field( + default=None, + metadata={ + "help": 'Maximum model sequence length for the colocated vLLM engine when `vllm_mode="colocate"`. Defaults to the model\'s maximum context length.' + }, + ) + vllm_model_impl: str = field( + default="vllm", + metadata={"help": 'Model implementation backend to use in vLLM. Use "vllm" (default) or "transformers".'}, + ) + vllm_structured_outputs_regex: str | None = field( + default=None, + metadata={"help": "Regex pattern used for vLLM structured outputs (optional)."}, + ) + vllm_sync_frequency: int = field( + default=1, + metadata={ + "help": "Frequency (in training steps) to synchronize model weights to the vLLM engine. Set to 1 to sync after every step." + }, + ) + vllm_enable_sleep_mode: bool = field( + default=False, + metadata={ + "help": "Enable vLLM sleep mode to offload student weights/cache during the optimizer step. Keeps GPU " + "memory usage low, but waking the engine adds host–device transfer latency." + }, + ) + # Parameters that control the logging + log_completions: bool = field( + default=False, + metadata={ + "help": "Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is " + "installed, it prints the sample. If `wandb` logging is enabled, it logs it to `wandb`." + }, + ) + log_completions_steps: int = field( + default=100, + metadata={ + "help": "Number of steps between logging (prompt, completion) pairs. Only used if `log_completions` is " + "set to `True`." + }, + ) + num_completions_to_print: int | None = field( + default=None, + metadata={"help": "Number of completions to print with `rich`. If `None`, all completions are logged."}, + ) + wandb_entity: str | None = field( + default=None, + metadata={"help": ("The entity to store runs under.")}, + ) + wandb_project: str | None = field( + default=None, + metadata={"help": ("The project to store runs under.")}, + ) + wandb_run_group: str | None = field( + default=None, + metadata={"help": ("The group to store runs under.")}, + ) + wandb_log_unique_prompts: bool = field( + default=True, + metadata={ + "help": ("Whether to log the unique prompts to wandb. This will create a new run for each unique prompt.") + }, + ) + callbacks: list[str] = field( + default_factory=lambda: [], + metadata={"help": "The callbacks to run during training."}, + ) + hub_model_revision: str | None = field( + default="main", metadata={"help": "The Hub model branch to push the model to."} + ) + overwrite_hub_revision: bool = field(default=False, metadata={"help": "Whether to overwrite the Hub revision."}) + push_to_hub_revision: bool = field(default=False, metadata={"help": "Whether to push to a Hub revision/branch."}) + + def __post_init__(self): + + # check lmbda and beta are in the range [0, 1] + if self.lmbda < 0.0 or self.lmbda > 1.0: + raise ValueError("lmbda must be in the range [0.0, 1.0].") + if self.beta < 0.0 or self.beta > 1.0: + raise ValueError("beta must be in the range [0.0, 1.0].") \ No newline at end of file