raulmodena/leire-corpus
Leire Corpus (EN) Tokenized pretraining corpus for Leire, a 343.7M-parameter Brazilian Portuguese LM trained from scratch on Kaggle T4s. ~15B tokens, 70% PT / 15% code / 8% math / 7% educational English, tokenized with a custom 32,768 BPE vocabulary trained on the same mixture. Shards are uint16 binaries; recipe and stats below (in Portuguese). Corpus de pre-treino da Leire, um LM de 343,7M de parametros em portugues brasileiro, treinado do zero em T4 do Kaggle. O projeto e… See the full description on the dataset page: https://huggingface.co/datasets/raulmodena/leire-corpus.
02.2k
1import json, statistics, sys, time2from pathlib import Path3import torch4RAIZ = Path('/kaggle/working')5sys.path.insert(0, str(RAIZ))6from leire.config import ModelConfig7from leire.model import LeireModel8from leire.optimizer import build_optimizers9 10SAIDA = RAIZ / 'relatorios' / 'throughput_proxy.json'11EST = {'medidas': {}}12 13 14def medir(cfg, mb, seq_len, modo, gc, passos=10, aquec=3):15 torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats()16 model = LeireModel(cfg).cuda(); model.recursion_mode = modo17 if gc:18 model.gradient_checkpointing_enable()19 model.train()20 muon, adamw, _ = build_optimizers(model)21 scaler = torch.amp.GradScaler('cuda', init_scale=2 ** 14)22 x = torch.randint(0, cfg.vocab_size, (mb, seq_len), device='cuda')23 y = torch.randint(0, cfg.vocab_size, (mb, seq_len), device='cuda')24 tempos = []25 for i in range(passos):26 torch.cuda.synchronize(); t0 = time.perf_counter()27 muon.zero_grad(set_to_none=True); adamw.zero_grad(set_to_none=True)28 with torch.autocast('cuda', dtype=torch.float16):29 perda = model(x, targets=y)['loss']30 scaler.scale(perda).backward()31 scaler.unscale_(muon); scaler.unscale_(adamw)32 torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)33 scaler.step(muon); scaler.step(adamw); scaler.update()34 torch.cuda.synchronize()35 if i >= aquec:36 tempos.append(time.perf_counter() - t0)37 pico = torch.cuda.max_memory_reserved() / 1024 ** 338 del model, muon, adamw; torch.cuda.empty_cache()39 med = statistics.median(tempos)40 return {'mb': mb, 'modo': modo, 'gc': gc, 'tok_s': round(mb * seq_len / med, 1),41 'pico_gb': round(pico, 2), 's_passo': round(med, 5)}42 43 44proxy = ModelConfig.proxy()45proxy_denso = ModelConfig.proxy_dense()46melhor = {}47for nome, cfg, modo in [('proxy_gather', proxy, 'gather'), ('proxy_dense', proxy, 'dense'),48 ('proxy_denso_baseline', proxy_denso, 'dense')]:49 for mb in [8, 16, 32, 48, 64]:50 try:51 r = medir(cfg, mb, 1024, modo, False)52 except (torch.cuda.OutOfMemoryError, RuntimeError) as e:53 if 'out of memory' not in str(e).lower():54 raise55 torch.cuda.empty_cache(); print(nome, 'mb', mb, 'OOM', flush=True); break56 print(nome, 'mb', mb, int(r['tok_s']), 'tok/s', r['pico_gb'], 'GB', flush=True)57 EST['medidas'][nome + '|mb' + str(mb)] = r58 SAIDA.write_text(json.dumps(EST, indent=1), encoding='utf-8')59 if r['pico_gb'] > 13.0:60 break61 melhor[nome] = max(melhor.get(nome, 0), r['tok_s'])62 63EST['melhor_por_config'] = melhor64if melhor.get('proxy_gather'):65 tps2 = melhor['proxy_gather'] * 1.866 for tokens in [200e6, 300e6, 500e6]:67 EST.setdefault('custo_ablacao', {})[str(int(tokens / 1e6)) + 'M_tokens'] = {68 'horas_1corrida_1gpu': round((tokens / melhor['proxy_gather']) / 3600, 2),69 'horas_par_pareado_2gpu': round((tokens / melhor['proxy_gather']) / 3600, 2),70 'horas_6_corridas': round(3 * (tokens / melhor['proxy_gather']) / 3600, 2)}71SAIDA.write_text(json.dumps(EST, indent=1), encoding='utf-8')72print('PROXY FIM', json.dumps(melhor), flush=True)73