train.py

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import pandas as pd
from transformers import AutoTokenizer, GPT2LMHeadModel, GPT2Config
from pathlib import Path
from tqdm import tqdm
import os
import re
import glob

torch.backends.cuda.matmul.allow_tf32 = True  # Для тензорных ядер Ampere+
torch.backends.cudnn.allow_tf32 = True

# Конфигурация
SEQ_LEN = 256
BATCH_SIZE = 32
EPOCHS = 100
LR = 0.0005
SPECIAL_PENALTY = 1.5
CHECKPOINT_DIR = "./checkpoints"
MAX_GRAD_NORM = 1.2
GRAD_ACCUM_STEPS = 4  # Подбирайте экспериментально
accum_step = 0  # Счётчик накопленных шагов
os.makedirs(CHECKPOINT_DIR, exist_ok=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.set_float32_matmul_precision('high')

# Инициализация токенизатора
tokenizer = AutoTokenizer.from_pretrained("sberbank-ai/rugpt3small_based_on_gpt2")
tokenizer.pad_token = tokenizer.eos_token

# Функция для определения спецтокенов
def get_special_token_ids(tokenizer):
    emoji_pattern = re.compile(
        "["
        "\U0001F600-\U0001F64F"
        "\U0001F300-\U0001F5FF"
        "\U0001F680-\U0001F6FF"
        "\U0001F1E0-\U0001F1FF"
        "\U00002500-\U00002BEF"
        "\U00002702-\U000027B0"
        "\U000024C2-\U0001F251"
        "]+", 
        flags=re.UNICODE
    )
    
    text_smiles = [
        ":-)", ":-(", ";-)", ":)", ":(", ":D", ":d",
        ":-D", ":-d", ":P", ":p", ";)", ":/", ":\\"
    ]
    
    special_token_ids = set()
    
    for token, token_id in tokenizer.vocab.items():
        if emoji_pattern.search(token):
            special_token_ids.add(token_id)
            
    for smile in text_smiles:
        tokens = tokenizer.encode(smile, add_special_tokens=False)
        special_token_ids.update(tokens)
    
    return special_token_ids

special_token_ids = get_special_token_ids(tokenizer)

# Кастомный лосс
class PriorityLoss(nn.Module):
    def __init__(self, special_ids, penalty=1.5, base_loss=nn.CrossEntropyLoss(ignore_index=-100)):
        super().__init__()
        self.register_buffer('special_ids', torch.tensor(list(special_ids), dtype=torch.long))  # Фикс 1
        self.penalty = penalty
        self.base_loss = base_loss

    def forward(self, logits, labels):
        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()
    
        # Перенос special_ids на устройство labels
        special_ids = self.special_ids.to(shift_labels.device)  # <-- Исправление
    
        base_loss = self.base_loss(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

        probs = torch.softmax(shift_logits, dim=-1)
        special_mask = torch.isin(shift_labels, special_ids)  # Используем локальный special_ids
    
        special_probs = probs[special_mask].sum()
        valid = (shift_labels != -100).sum()
        penalty = special_probs / valid if valid > 0 else 0.0

        return base_loss + self.penalty * penalty

criterion = PriorityLoss(special_token_ids, penalty=SPECIAL_PENALTY)

# Простой датасет без фильтрации
class TextDataset(Dataset):
    def __init__(self, df, tokenizer, max_length=256):
        self.texts = df['text'].tolist()
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        encoding = self.tokenizer(
            self.texts[idx],
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        return {
            'input_ids': encoding['input_ids'].squeeze(),
            'attention_mask': encoding['attention_mask'].squeeze()
        }

# Инициализация модели GPT
config = GPT2Config(
    vocab_size=tokenizer.vocab_size,
    n_positions=SEQ_LEN,
    n_embd=512,
    n_layer=6,
    n_head=8,
    pad_token_id=tokenizer.pad_token_id
)
model = GPT2LMHeadModel(config).to(device)
model = torch.compile(model)

optimizer = optim.AdamW(model.parameters(), lr=LR, weight_decay=0.03, fused=True)
scaler = torch.amp.GradScaler()

# Загрузка данных
train_df = pd.read_parquet("train.parquet")
test_df = pd.read_parquet("test.parquet")
train_dataset = TextDataset(train_df, tokenizer, SEQ_LEN)
test_dataset = TextDataset(test_df, tokenizer, SEQ_LEN)
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)

# Функции для работы с чекпоинтами
def find_latest_checkpoint():
    checkpoints = glob.glob(f"{CHECKPOINT_DIR}/checkpoint_epoch*_step*.pt")
    if not checkpoints:
        return None, 0, 0
    
    max_epoch = -1
    max_step = -1
    best_cp = None
    
    for cp in checkpoints:
        filename = os.path.basename(cp)
        match = re.match(r"checkpoint_epoch(\d+)_step(\d+)\.pt", filename)
        if match:
            epoch = int(match.group(1))
            step = int(match.group(2))
            if (epoch > max_epoch) or (epoch == max_epoch and step > max_step):
                max_epoch = epoch
                max_step = step
                best_cp = cp
    
    return best_cp, max_epoch, max_step

# Загрузка чекпоинта
latest_checkpoint, loaded_epoch, loaded_step = find_latest_checkpoint()
initial_epoch = 0
global_step = 0

if latest_checkpoint:
    print(f"Loading checkpoint: {latest_checkpoint}")
    checkpoint = torch.load(latest_checkpoint, map_location=device)
    
    model.load_state_dict(checkpoint['model'])
    optimizer.load_state_dict(checkpoint['optimizer'])
    scaler.load_state_dict(checkpoint['scaler'])
    
    initial_epoch = loaded_epoch
    global_step = checkpoint['step']
    
    print(f"Resumed training from epoch {initial_epoch}, step {global_step}")

# Генерация текста
def balanced_generate(model, tokenizer, prompt, max_len=100, temp=0.9, top_p=0.95, 
                     special_penalty=2.0, min_penalty=0.5):
    model.eval()
    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
    special_tensor = torch.tensor(list(special_token_ids)).to(device)
    
    for _ in range(max_len):
        with torch.no_grad():
            outputs = model(input_ids)
            logits = outputs.logits[:, -1, :] / temp
            
            seq_len = input_ids.size(1)
            penalty = special_penalty * (1 - seq_len/max_len) + min_penalty
            logits[:, special_tensor] -= penalty
            
            sorted_logits, sorted_indices = torch.sort(logits, descending=True)
            cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
            
            sorted_indices_to_remove = cumulative_probs > top_p
            sorted_indices_to_remove[..., 0] = False
            
            indices_to_remove = sorted_indices_to_remove.scatter(
                1, sorted_indices, sorted_indices_to_remove
            )
            logits = logits.masked_fill(indices_to_remove, float('-inf'))
            
            probs = torch.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            
            input_ids = torch.cat([input_ids, next_token], dim=1)
            
    return tokenizer.decode(input_ids[0], skip_special_tokens=True)

# Тренировочный цикл
for epoch in range(initial_epoch, EPOCHS):
    model.train()
    optimizer.zero_grad()  # Обнуление градиентов в начале эпохи
    
    # Сбрасываем счётчик накопления для каждой эпохи
    accum_step = 0
    
    for batch_idx, batch in enumerate(tqdm(train_loader)):
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        
        labels = input_ids.clone()
        labels[attention_mask == 0] = -100

        with torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):
            outputs = model(input_ids, attention_mask=attention_mask)
            loss = criterion(outputs.logits, labels)
        
        # Нормализация лосса для накопления
        loss = loss / GRAD_ACCUM_STEPS
        
        # Backward с масштабированием
        scaler.scale(loss).backward()
        accum_step += 1

        # Шаг только при накоплении достаточного числа градиентов
        if accum_step % GRAD_ACCUM_STEPS == 0 or batch_idx == len(train_loader)-1:
            # Обрезка градиентов перед шагом
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), MAX_GRAD_NORM)
            
            # Обновление весов и масштаба
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()
            
            # Логирование и чекпоинты
            current_loss = loss.item() * GRAD_ACCUM_STEPS  # Восстанавливаем оригинальный лосс
        if global_step % 100 == 0:
            print(f"Epoch {epoch} | Step {global_step} | Loss: {loss.item():.4f}")

        if global_step % 1000 == 0:
            model.eval()
            with torch.inference_mode():
                text = balanced_generate(
                    model,
                    tokenizer,
                    prompt="Я думаю",
                    special_penalty=3.0,
                    min_penalty=0.7
                )
                print(f"\nGenerated:\n{text}\n")
                
                torch.save({
                    'model': model.state_dict(),
                    'optimizer': optimizer.state_dict(),
                    'scaler': scaler.state_dict(),
                    'step': global_step
                }, f"{CHECKPOINT_DIR}/checkpoint_epoch{epoch}_step{global_step}.pt")
            
            model.train()

        global_step += 1