import json import logging import math import os import time from functools import partial from pathlib import Path import torch import torch.nn.functional as F from huggingface_hub import HfApi from safetensors import safe_open from torch import nn from torch.nn.attention.flex_attention import BlockMask, create_block_mask, flex_attention from torch.utils.tensorboard import SummaryWriter from tqdm.auto import tqdm LOGGER = logging.getLogger(__name__) HF_TOKEN = os.environ['HF_TOKEN'] API = HfApi(token=HF_TOKEN) class FixedPositionalEncoding(nn.Module): def __init__(self, hidden_size, max_sequence_length=512): super().__init__() self.hidden_size = hidden_size self.max_sequence_length = max_sequence_length pos_enc = torch.zeros(max_sequence_length, hidden_size) position = torch.arange(0.0, max_sequence_length).unsqueeze(1) coef = -math.log(10000.0) / hidden_size div_term = torch.exp(coef * torch.arange(0.0, hidden_size, 2)) pos_enc[:, 0::2] = torch.sin(position * div_term) pos_enc[:, 1::2] = torch.cos(position * div_term) pos_enc.div_(math.sqrt(hidden_size)) self.register_buffer('pos_enc', pos_enc) def forward(self, position_ids): return torch.index_select(self.pos_enc, 0, position_ids.reshape(-1)).reshape(*position_ids.shape, -1) class DecoderAttention(nn.Module): def __init__(self, hidden_size, num_heads, layer_idx, kv=True): super().__init__() self.hidden_size = hidden_size self.num_heads = num_heads self.layer_idx = layer_idx self.head_dim = hidden_size // num_heads self.scale = self.head_dim**-0.5 self.query_net = nn.Linear(hidden_size, hidden_size) if kv: self.key_net = nn.Linear(hidden_size, hidden_size) self.value_net = nn.Linear(hidden_size, hidden_size) self.out_projection = nn.Linear(hidden_size, hidden_size) def _reshape(self, x): b, t, _ = x.shape return x.view(b, t, self.num_heads, self.head_dim).transpose(1, 2) def forward( self, hidden_states, context_states=None, attention_mask=None, past_key_values=None, diffusion=False, ): self_attn = context_states is None name = 'self' if self_attn else 'cross' if self_attn: context_states = hidden_states bsz, tgt_len, _ = hidden_states.size() query = self._reshape(self.query_net(hidden_states)) if diffusion: ar_key = past_key_values[f'{self.layer_idx}.{name}.key'] ar_value = past_key_values[f'{self.layer_idx}.{name}.value'] if self_attn: diff_key = self._reshape(self.key_net(context_states)) diff_value = self._reshape(self.value_net(context_states)) key = torch.cat((ar_key, diff_key), dim=2) value = torch.cat((ar_value, diff_value), dim=2) else: key = ar_key value = ar_value else: key = self._reshape(self.key_net(context_states)) value = self._reshape(self.value_net(context_states)) past_key_values[f'{self.layer_idx}.{name}.key'] = key past_key_values[f'{self.layer_idx}.{name}.value'] = value attn_output = flex_attention(query, key, value, block_mask=attention_mask, scale=self.scale) attn_output = attn_output.transpose(1, 2).reshape(bsz, tgt_len, self.hidden_size) return self.out_projection(attn_output) class DecoderFeedForward(nn.Module): def __init__(self, hidden_size, inner_size, hidden_act='relu'): super().__init__() self.dense_in = nn.Linear(hidden_size, inner_size) assert hidden_act == 'relu' self.activation = nn.ReLU() self.dense_out = nn.Linear(inner_size, hidden_size) def forward(self, x): return self.dense_out(self.activation(self.dense_in(x))) class TransformerDecoderLayer(nn.Module): def __init__(self, hidden_size, inner_size, num_heads, diffusion, layer_idx, hidden_act='relu'): super().__init__() self.layer_norm_1 = nn.LayerNorm(hidden_size) self.first_sub_layer = DecoderAttention(hidden_size, num_heads, layer_idx=layer_idx) self.layer_norm_2 = nn.LayerNorm(hidden_size) self.second_sub_layer = DecoderAttention(hidden_size, num_heads, layer_idx=layer_idx, kv=not diffusion) self.layer_norm_3 = nn.LayerNorm(hidden_size) self.third_sub_layer = DecoderFeedForward(hidden_size, inner_size, hidden_act=hidden_act) def forward( self, hidden_states, encoder_hidden_states=None, self_attention_mask=None, cross_attention_mask=None, past_key_values=None, diffusion=False, ): residual = hidden_states hidden_states = self.layer_norm_1(hidden_states) self_out = self.first_sub_layer( hidden_states, context_states=None, attention_mask=self_attention_mask, past_key_values=past_key_values, diffusion=diffusion, ) hidden_states = residual + self_out residual = hidden_states hidden_states = self.layer_norm_2(hidden_states) cross_out = self.second_sub_layer( hidden_states, context_states=encoder_hidden_states, attention_mask=cross_attention_mask, past_key_values=past_key_values, diffusion=diffusion, ) hidden_states = residual + cross_out residual = hidden_states hidden_states = self.layer_norm_3(hidden_states) hidden_states = residual + self.third_sub_layer(hidden_states) return hidden_states class TransformerDecoderEmbedding(nn.Module): def __init__(self, vocab_size, hidden_size, max_sequence_length, padding_idx=2): super().__init__() self.token_embedding = nn.Embedding(vocab_size, hidden_size, padding_idx) self.position_embedding = FixedPositionalEncoding(hidden_size, max_sequence_length) self.layer_norm = nn.LayerNorm(hidden_size) def forward(self, input_ids, positions): token_embeds = self.token_embedding(input_ids) if positions is None: pos_embeds = self.position_embedding.pos_enc[:input_ids.shape[-1]] else: pos_embeds = self.position_embedding(positions) return self.layer_norm(token_embeds + pos_embeds) class TransformerDecoderCore(nn.Module): def __init__(self, hidden_size, inner_size, num_heads, num_layers, diffusion, hidden_act='relu'): super().__init__() self.layers = nn.ModuleList( [ TransformerDecoderLayer(hidden_size, inner_size, num_heads, diffusion, layer_idx=i, hidden_act=hidden_act) for i in range(num_layers) ] ) self.final_layer_norm = nn.LayerNorm(hidden_size) def forward( self, hidden_states, encoder_hidden_states=None, self_attention_mask=None, cross_attention_mask=None, past_key_values=None, diffusion=False ): for layer in self.layers: hidden_states = layer( hidden_states, encoder_hidden_states=encoder_hidden_states, self_attention_mask=self_attention_mask, cross_attention_mask=cross_attention_mask, past_key_values=past_key_values, diffusion=diffusion, ) return self.final_layer_norm(hidden_states) class TransformerDecoderWrapper(nn.Module): def __init__(self, config, diffusion=False): super().__init__() if not isinstance(config, dict): with open(config, 'r', encoding='utf-8') as file: config = json.load(file) dec_config = config['transf_decoder']['config_dict'] hidden_size = dec_config['hidden_size'] self._embedding = TransformerDecoderEmbedding( vocab_size=config['vocab_size'], hidden_size=hidden_size, max_sequence_length=dec_config['max_sequence_length'], padding_idx=2, ) self._decoder = TransformerDecoderCore( hidden_size=hidden_size, inner_size=dec_config['inner_size'], num_heads=dec_config['num_attention_heads'], num_layers=dec_config['num_layers'], diffusion=diffusion, hidden_act=dec_config.get('hidden_act', 'relu'), ) self.diffusion = diffusion self._lm_head = nn.Linear(hidden_size, config['vocab_size']) self._lm_head.weight = self._embedding.token_embedding.weight def load(self, model: str): state = {} with safe_open(model, 'pt') as file: for key in file.keys(): if self.diffusion and ('.second_sub_layer.key_net.' in key or '.second_sub_layer.value_net.' in key): continue elif key.startswith('transf_decoder.'): state[key.removeprefix('transf_decoder.')] = file.get_tensor(key) elif key == 'log_softmax.mlp.layer0.bias': state['_lm_head.bias'] = file.get_tensor(key) elif key == 'log_softmax.mlp.layer0.weight': state['_lm_head.weight'] = file.get_tensor(key) self.load_state_dict(state) return self def forward( self, input_ids, positions=None, encoder_hidden_states=None, self_attention_mask=None, cross_attention_mask=None, past_key_values=None, diffusion=False ): hidden_states = self._embedding(input_ids, positions) hidden_states = self._decoder( hidden_states, encoder_hidden_states=encoder_hidden_states, self_attention_mask=self_attention_mask, cross_attention_mask=cross_attention_mask, past_key_values=past_key_values, diffusion=diffusion, ) return self._lm_head(hidden_states) def cosine_schedule(step: int, *, warmup_steps: int, max_steps: int) -> float: if step < warmup_steps: return step / warmup_steps progress = (step - warmup_steps) / (max_steps - warmup_steps) return 0.5 * (1 + math.cos(math.pi * progress)) @torch.no_grad() def load(model_id: str, hub_id: str, warmup_steps: int, max_steps: int, output_dir: Path) -> tuple[ int, TransformerDecoderWrapper, TransformerDecoderCore, torch.optim.AdamW, torch.optim.lr_scheduler.LambdaLR, ]: config = API.hf_hub_download(model_id, 'config.json') teacher = TransformerDecoderWrapper(config).to(device='cuda', dtype=torch.bfloat16).eval() student = TransformerDecoderWrapper(config, diffusion=True).to(device='cuda', dtype=torch.bfloat16) model = API.hf_hub_download(model_id, 'model.safetensors') teacher.load(model) student.load(model) student = student._decoder teacher.requires_grad_(False) cosine_lr = partial(cosine_schedule, warmup_steps=warmup_steps, max_steps=max_steps) decay_params = [] no_decay_params = [] for param in student.parameters(): if param.ndim < 2: no_decay_params.append(param) else: decay_params.append(param) optimizer = torch.optim.AdamW([ {'params': decay_params, 'weight_decay': 0.01}, {'params': no_decay_params, 'weight_decay': 0.0}, ], lr=4e-3, betas=(0.9, 0.98), eps=1e-8) scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, cosine_lr) cur_step = 0 checkpoints = [x.path for x in API.list_bucket_tree(hub_id) if x.path.endswith('.pt')] checkpoints.sort(key=lambda x: int(x.split('.')[0])) if checkpoints: checkpoint = checkpoints[-1] LOGGER.info('Checkpoint found %s', checkpoint) API.download_bucket_files(hub_id, files=[(checkpoint, output_dir/checkpoint)]) state = torch.load(output_dir/checkpoint, weights_only=True) cur_step = state['step'] student.load_state_dict(state['model']) optimizer.load_state_dict(state['optimizer']) scheduler.load_state_dict(state['scheduler']) else: LOGGER.info('Checkpoint not detected') return cur_step, teacher, student, optimizer, scheduler @torch.compile(mode='default', fullgraph=True, dynamic=False) def forward( teacher: TransformerDecoderWrapper, student: TransformerDecoderCore, encoder_states: torch.Tensor, # [b,t,d] input_ids: torch.LongTensor, # [b,s] indices: torch.LongTensor, # [b,u] causal_attn_mask: BlockMask, self_attn_mask: BlockMask, cross_attn_mask: BlockMask, bsz: int, ar_len: int, diff_len: int, kv_len: int, diff_block: int, diff_blocks: int, pad_id: int, eot_id: int, ) -> torch.Tensor: past_key_values = {} with torch.no_grad(): teacher_attn_mask = cross_attn_mask._adjust(ar_len, kv_len) teacher_attn_mask.seq_lengths = ar_len, kv_len ar_logits = teacher( input_ids=input_ids, positions=None, encoder_hidden_states=encoder_states, self_attention_mask=causal_attn_mask, cross_attention_mask=teacher_attn_mask, past_key_values=past_key_values, diffusion=False, ) ar_logprobs = torch.log_softmax(ar_logits, dim=2, dtype=torch.float32) positions = (indices[:, :, None] + torch.arange(0, diff_block, device=input_ids.device)).view(bsz, diff_len) valid = positions < ar_len pad_mask = valid & (input_ids.gather(1, torch.where(valid, positions, 0)) > eot_id) target = ar_logprobs.take_along_dim(torch.clamp_max(positions, ar_len-1)[:, :, None], dim=1) del ar_logits, ar_logprobs train_ids = input_ids.gather(1, indices)[:, :, None] train_ids = F.pad(train_ids, (0, diff_block - 1), 'constant', pad_id).reshape(bsz, diff_len) diff_states = teacher._embedding(train_ids, positions) diff_states = student( diff_states, encoder_hidden_states=encoder_states, self_attention_mask=self_attn_mask, cross_attention_mask=cross_attn_mask, past_key_values=past_key_values, diffusion=True, ) diff_logits = teacher._lm_head(diff_states) diff_logprobs = torch.log_softmax(diff_logits, dim=2, dtype=torch.float32) kld = (target.exp() * (target - diff_logprobs)).sum(dim=2) # [b,s] weight = torch.exp(torch.arange(diff_block, dtype=torch.float32, device=kld.device) / -12) kld.view(bsz, diff_blocks, diff_block).mul_(weight) return (kld * pad_mask).sum() / pad_mask.sum() def train_step( encoder_states: torch.Tensor, # [b,t,d] encoder_lengths: torch.LongTensor, # [b] input_ids: torch.LongTensor, # [b,s] teacher: TransformerDecoderWrapper, student: TransformerDecoderCore, optimizer: torch.optim.AdamW, scheduler: torch.optim.lr_scheduler.LambdaLR, ) -> dict[str, torch.Tensor]: device = encoder_states.device bsz = encoder_states.shape[0] ar_len = input_ids.shape[1] diff_len = 2048 kv_len = encoder_states.shape[1] diff_block = 32 ar_blocks = ar_len // diff_block diff_blocks = diff_len // diff_block prompt_len = 9 pad_id = 2 eot_id = 3 specials = 254 input_lengths = (input_ids > eot_id).cumprod(dim=1).sum(dim=1) weights = (input_ids > specials).bfloat16() weights[:, prompt_len] = 1 indices = torch.multinomial(weights, diff_blocks).sort(dim=1).values _create_block_mask = torch.compile(create_block_mask, fullgraph=True, dynamic=False) def causal_mask(b, h, q_idx, kv_idx): return kv_idx <= q_idx causal_attn_mask = _create_block_mask(causal_mask, B=None, H=None, Q_LEN=ar_len, KV_LEN=ar_len, device=device) def self_mask(b, h, q_idx, kv_idx): block = q_idx // diff_block == kv_idx // diff_block - ar_blocks idx = q_idx // diff_block prefix = kv_idx < indices[b, idx] return (block | prefix) & (idx < input_lengths[b]) self_attn_mask = _create_block_mask(self_mask, B=bsz, H=None, Q_LEN=diff_len, KV_LEN=ar_len+diff_len, device=device) def cross_mask(b, h, q_idx, kv_idx): return kv_idx < encoder_lengths[b] cross_attn_mask = _create_block_mask(cross_mask, B=bsz, H=None, Q_LEN=diff_len, KV_LEN=kv_len, device=device) optimizer.zero_grad() loss = forward( teacher=teacher, student=student, encoder_states=encoder_states, input_ids=input_ids, indices=indices, causal_attn_mask=causal_attn_mask, self_attn_mask=self_attn_mask, cross_attn_mask=cross_attn_mask, bsz=bsz, ar_len=ar_len, diff_len=diff_len, kv_len=kv_len, diff_block=diff_block, diff_blocks=diff_blocks, pad_id=pad_id, eot_id=eot_id, ) loss.backward() grad_norm = nn.utils.clip_grad_norm_(student.parameters(), 1.0) optimizer.step() scheduler.step() return { 'train/loss': loss.detach().clone(), 'train/grad_norm': grad_norm.detach().clone(), 'train/learning_rate': scheduler.get_last_lr()[0], } def main(): model_id = 'efwkjn/cohere-asr-ja' hub_id = 'efwkjn/checkpoints' output_dir = Path('checkpoints') summary_writer = SummaryWriter(log_dir=output_dir/'runs') handler = logging.StreamHandler() formatter = logging.Formatter('%(levelname)s: %(message)s') handler.setFormatter(formatter) LOGGER.addHandler(handler) LOGGER.setLevel(logging.INFO) API.create_bucket(hub_id, private=True, exist_ok=True) max_steps = 2**17 warmup_steps = 2**13 save_steps = 2**10 logging_steps = 2**4 cur_step, teacher, student, optimizer, scheduler = load(model_id, hub_id, warmup_steps, max_steps, output_dir) dataset = None train_metrics: list[dict[str, torch.Tensor | float]] = [] time_start = time.perf_counter() for batch in tqdm(dataset, initial=cur_step, total=max_steps): cur_step += 1 metrics = train_step( **batch, teacher=teacher, student=student, optimizer=optimizer, scheduler=scheduler, ) if cur_step % logging_steps == 0: metrics['step'] = cur_step train_metrics.append(metrics) if cur_step % (logging_steps * 16) == 0: prev_metrics = train_metrics[-2] s = ' | '.join(f'{k[6:]}: {v.item():.5f}' for k, v in prev_metrics.items() if k.startswith('train/') and isinstance(v, torch.Tensor)) LOGGER.info(f'{prev_metrics["step"]}: {s}') if cur_step % save_steps == 0: train_time = time.perf_counter() - time_start summary_writer.add_scalar('train/time', train_time, cur_step) for m in train_metrics: step = m.pop('step') for k, v in m.items(): summary_writer.add_scalar(k, v, step) summary_writer.flush() train_metrics = [] time_start = time.perf_counter() checkpoint = output_dir/f'{cur_step}.pt' torch.save({ 'step': cur_step, 'model': student.state_dict(), 'optimizer': optimizer.state_dict(), 'scheduler': scheduler.state_dict(), }, checkpoint) API.run_as_future(API.sync_bucket, str(output_dir), f'hf://buckets/{hub_id}', ignore_times=True) if __name__ == '__main__': main()