train.py

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()