Release best Gemini-tuned Whisper Base and Small with code, normalization and evaluation
cd9b2d8 verified | #!/usr/bin/env python3 | |
| """Bounded distributed full-FT and real P3 evaluation compatibility check.""" | |
| import os | |
| import sys | |
| import torch | |
| from torch.nn.parallel import DistributedDataParallel | |
| from torch.utils.data import DataLoader, Subset | |
| from study_paths import CODE, GEMINI, ROOT, RELEASE, read, write | |
| sys.path.insert(0, str(CODE / 'gemini_finetune')) | |
| from data import GeminiDataset, GeminiCollator | |
| sys.path.insert(0, str(CODE)) | |
| from layered_multitask_model import LayeredMultiTaskWhisper | |
| from train_layered_curriculum import gpu_batch, loss_terms | |
| def main(): | |
| rank = int(os.environ['LOCAL_RANK']) | |
| device = torch.device('cuda', rank) | |
| torch.cuda.set_device(device) | |
| torch.set_num_threads(2) | |
| torch.distributed.init_process_group('nccl') | |
| canary = GeminiDataset(GEMINI / 'parser_canary_complete', 'train', require_complete=False) | |
| classes = read(RELEASE / 'classes.json')['names'] | |
| cfg = read(CODE / 'gemini_finetune/configs/whisper_base.json') | |
| core = LayeredMultiTaskWhisper(cfg['model_config'], len(classes), initialize_pretrained=False) | |
| core.load_state_dict(torch.load(cfg['checkpoint'], map_location='cpu', weights_only=False)['model'], strict=True) | |
| core.to(device).train().requires_grad_(True) | |
| model = DistributedDataParallel(core, device_ids=[rank], find_unused_parameters=False) | |
| optimizer = torch.optim.AdamW([ | |
| {'params': list(core.encoder.parameters()), 'lr': cfg['encoder_lr']}, | |
| {'params': [p for n,p in core.named_parameters() if not n.startswith('encoder.')], 'lr': cfg['head_lr']}]) | |
| data = Subset(canary, [(rank + step * 4) % len(canary) for step in range(2)]) | |
| loader = DataLoader(data, batch_size=1, num_workers=1, multiprocessing_context='spawn', | |
| collate_fn=GeminiCollator(cfg['model_config'], canary.norm)) | |
| weights = torch.ones(len(classes), device=device) | |
| for raw in loader: | |
| batch = gpu_batch(raw, device) | |
| optimizer.zero_grad(set_to_none=True) | |
| with torch.autocast('cuda', torch.bfloat16): | |
| output = model(batch['mel'], batch['mel_mask'], batch['event_starts'], batch['event_ends']) | |
| loss = sum(loss_terms(output, batch, weights).values()) | |
| loss.backward() | |
| assert torch.isfinite(loss) and all(p.grad is None or torch.isfinite(p.grad).all() for p in core.parameters()) | |
| optimizer.step() | |
| torch.distributed.barrier() | |
| del model, core, optimizer, output, loss, batch, loader | |
| torch.cuda.empty_cache() | |
| # Check the new P3 wrapper, spawned audio loader and predicted event path. | |
| local_result = {'rank': rank, 'ddp_iterations': 2, 'finite_loss_and_gradients': True} | |
| if rank < 2: | |
| from evaluate_whisper import P3WhisperData, HoldoutLoader, forward | |
| from train_probes import evaluate | |
| size = ('base', 'small')[rank] | |
| cfg = read(CODE / 'gemini_finetune/configs' / ('whisper_' + size + '.json')) | |
| core = LayeredMultiTaskWhisper(cfg['model_config'], len(classes), initialize_pretrained=False).to(device) | |
| core.load_state_dict(torch.load(cfg['checkpoint'], map_location='cpu', weights_only=False)['model'], strict=True) | |
| checked = {} | |
| for split in ('validation', 'test'): | |
| metrics = evaluate(core, Subset(P3WhisperData(split), [0]), device, weights, True, | |
| HoldoutLoader(RELEASE / size, canary.norm), forward) | |
| assert metrics['clips'] == 1 and len(metrics['per_score']) == 192 | |
| checked[split] = {'clips': 1, 'targets': 192, 'predicted_and_oracle_burst_paths': True} | |
| local_result.update(model=size, p3_checks=checked) | |
| results = [None] * 4 | |
| torch.distributed.all_gather_object(results, local_result) | |
| if rank == 0: | |
| write(ROOT / 'DISTRIBUTED_PIPELINE_CONTRACT.json', { | |
| 'passed': True, 'scope': 'Two real canary updates per rank and one P3 clip per split/model; software check, not an evaluation result', | |
| 'ranks': results}) | |
| print('DISTRIBUTED_PIPELINE_CONTRACT_PASS', flush=True) | |
| torch.distributed.destroy_process_group() | |
| if __name__ == '__main__': | |
| main() | |