pretrain.py

import os
import yaml
import argparse
import torch
import torch.optim as optim
import torch.distributed as dist
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
from tqdm import tqdm

from asr.data import AudioDataset, AudioCollator
from asr.model import (
    CNNFeatureExtractor,
    ConvPositionalEncoding,
    GumbelVectorQuantizer,
    TransformerEncoder,
)
from asr.loss import InfoNCEContrastiveLoss
from asr.utils import compute_mask_indices, apply_mask, NegativeSampler

def setup_ddp():
    """Initializes the distributed process group."""
    dist.init_process_group("nccl")
    local_rank = int(os.environ['LOCAL_RANK'])
    torch.cuda.set_device(local_rank)
    device = torch.device(f'cuda:{local_rank}')
    rank = int(os.environ['RANK'])
    world_size = int(os.environ['WORLD_SIZE'])
    return rank, world_size, device

def cleanup_ddp():
    """Cleans up the distributed process group."""
    dist.destroy_process_group()

def main(rank, world_size, device, config):
    if rank == 0:
        print(f"Distributed training on {world_size} GPUs")
        os.makedirs(config['checkpoint_dir'], exist_ok=True)

    # 1. Data Loading
    dataset = AudioDataset(manifest_path=config['manifest_path'])
    collator = AudioCollator()
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    data_loader = DataLoader(
        dataset,
        batch_size=config['batch_size'],
        collate_fn=collator,
        num_workers=4,
        pin_memory=True,
        sampler=sampler
    )

    # 2. Model Instantiation
    if rank == 0: print("Instantiating models...")
    cnn_feat = DDP(CNNFeatureExtractor().to(device), device_ids=[device.index])
    post_extract_proj = DDP(nn.Linear(512, config['embed_dim']).to(device), device_ids=[device.index])
    quantizer = DDP(GumbelVectorQuantizer(
        dim=config['embed_dim'], num_vars=320, temp=(2.0, 0.5, 0.999995), groups=2
    ).to(device), device_ids=[device.index])
    encoder = DDP(TransformerEncoder(
        embed_dim=config['embed_dim'], ffn_dim=config['ffn_dim'], num_heads=config['num_heads'],
        num_layers=config['num_layers'], dropout=config['dropout']
    ).to(device), device_ids=[device.index])
    pos_encoder = DDP(ConvPositionalEncoding(embed_dim=config['embed_dim']).to(device), device_ids=[device.index])
    mask_embedding = nn.Parameter(torch.FloatTensor(config['embed_dim']).uniform_().to(device))

    # 3. Loss, Sampler, Optimizer, and Scheduler
    criterion = InfoNCEContrastiveLoss(temperature=0.1).to(device)
    negative_sampler = NegativeSampler(n_negatives=100).to(device)

    all_params = (
        list(cnn_feat.module.parameters()) + list(post_extract_proj.module.parameters()) +
        list(quantizer.module.parameters()) + list(encoder.module.parameters()) +
        list(pos_encoder.module.parameters()) + [mask_embedding]
    )
    optimizer = optim.AdamW(all_params, lr=config['learning_rate'], betas=(0.9, 0.98), eps=1e-6)

    total_steps = len(data_loader) * config['num_epochs']
    warmup_steps = int(total_steps * config['warmup_ratio'])
    
    def scheduler_fn(step):
        if step < warmup_steps: return float(step) / float(max(1, warmup_steps))
        return max(0.0, float(total_steps - step) / float(max(1, total_steps - warmup_steps)))
    scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=scheduler_fn)
    
    # 4. Training Loop
    if rank == 0: print("--- Starting Training ---")
    global_step = 0
    for epoch in range(config['num_epochs']):
        sampler.set_epoch(epoch)
        progress_bar = tqdm(data_loader, desc=f"Epoch {epoch+1}", disable=(rank != 0))
        
        for batch_waveforms, padding_mask in progress_bar:
            if batch_waveforms is None: continue
            batch_waveforms, padding_mask = batch_waveforms.to(device), padding_mask.to(device)
            
            optimizer.zero_grad(set_to_none=True)
            
            feats = cnn_feat(batch_waveforms).permute(0, 2, 1)
            proj = post_extract_proj(feats)
            unmasked_features = proj.clone()
            
            proj_with_pos = pos_encoder(proj)
            mask_indices_np = compute_mask_indices(proj.shape[:2], config['mask_prob'], config['mask_length'])
            mask_indices = torch.from_numpy(mask_indices_np).to(device) & ~padding_mask
            masked_proj = apply_mask(proj_with_pos, mask_indices, mask_embedding)
            
            quantizer.module.set_num_updates(global_step)
            quant_x, diversity_loss = quantizer(unmasked_features)
            
            ctx = encoder(masked_proj, padding_mask=padding_mask)
            masked_ctx, quant_tgt = ctx[mask_indices], quant_x[mask_indices]
            
            if masked_ctx.size(0) == 0: continue
            
            negs = negative_sampler.sample_negatives(quant_x, num_masked=masked_ctx.size(0), padding_mask=padding_mask)
            contrastive_loss, _ = criterion(masked_ctx, quant_tgt, negs)
            loss = contrastive_loss + config['diversity_loss_weight'] * diversity_loss.sum()
            
            if torch.isinf(loss) or torch.isnan(loss):
                if rank == 0: print(f"Warning: Skipping step {global_step} due to inf/nan loss.")
                continue
            
            loss.backward()
            torch.nn.utils.clip_grad_norm_(all_params, 1.0)
            optimizer.step()
            scheduler.step()
            
            if rank == 0: progress_bar.set_postfix(loss=f"{loss.item():.4f}")
            global_step += 1
        
        # 5. Checkpointing
        if rank == 0:
            print("Epoch finished. Saving checkpoint...")
            checkpoint_path = os.path.join(config['checkpoint_dir'], f"checkpoint_epoch_{epoch+1}.pt")
            torch.save({
                'epoch': epoch, 'global_step': global_step,
                'cnn_feat_state_dict': cnn_feat.module.state_dict(),
                'post_extract_proj_state_dict': post_extract_proj.module.state_dict(),
                'quantizer_state_dict': quantizer.module.state_dict(),
                'encoder_state_dict': encoder.module.state_dict(),
                'pos_encoder_state_dict': pos_encoder.module.state_dict(),
                'optimizer_state_dict': optimizer.state_dict()
            }, checkpoint_path)
            print(f"Checkpoint saved to {checkpoint_path}")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Pre-train a SSL ASR model.")
    parser.add_argument('--config', type=str, required=True, help='Path to the YAML configuration file.')
    args = parser.parse_args()

    with open(args.config, 'r') as f:
        config = yaml.safe_load(f)

    rank, world_size, device = setup_ddp()
    main(rank, world_size, device, config)
    cleanup_ddp()