diff --git a/scripts/10_train.py b/scripts/10_train.py index f750e7a..382d669 100644 --- a/scripts/10_train.py +++ b/scripts/10_train.py @@ -67,7 +67,13 @@ from pathlib import Path import torch from datasets import Dataset from peft import LoraConfig, get_peft_model -from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments +from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + Trainer, + TrainerCallback, + TrainingArguments, +) REPO_ROOT = Path(__file__).resolve().parent.parent @@ -95,6 +101,16 @@ LORA_DROPOUT = float(os.environ.get("LORA_DROPOUT", "0.05")) EVAL_STEPS = int(os.environ.get("EVAL_STEPS", "50")) SAVE_STEPS = int(os.environ.get("SAVE_STEPS", "50")) GRAD_ACCUM = int(os.environ.get("GRAD_ACCUM", "16")) +# El default de HF Trainer para per_device_eval_batch_size es 8, INDEPENDIENTE del batch de +# entrenamiento. Nunca se fijaba explicitamente. Con evaluation_loop corriendo sin gradient +# checkpointing (no hace falta, no hay backward) y sin el truco de recompute, un lote de eval de +# 8 secuencias largas se materializa entero -- logits de forma (8, seq_len, vocab_size) en fp32 +# via convert_to_fp32 -- y con vocab grande eso son decenas de GB en una sola asignacion. Es la +# causa real de los dos OOM de esta fase (el smoke run del 32-mas-largos Y la calibracion de 8 +# ejemplos): en ambos casos NO fue el forward/backward de entrenamiento (medido establemente en +# ~74GB para cualquier longitud de 2808 a 3265 tokens), fue el forward de evaluacion agrupando +# examples largos en un batch de 8. +EVAL_BATCH_SIZE = int(os.environ.get("EVAL_BATCH_SIZE", "1")) MAX_TOKENS = int(os.environ["MAX_TOKENS"]) if os.environ.get("MAX_TOKENS") else None PRESERVE_THINKING = os.environ.get("PRESERVE_THINKING", "").lower() in ("1", "true", "yes") ALLOW_OVERWRITE = os.environ.get("ALLOW_OVERWRITE", "").lower() in ("1", "true", "yes") @@ -129,6 +145,7 @@ def print_banner(args): ("learning_rate", LEARNING_RATE), ("num_train_epochs", NUM_EPOCHS), ("grad_accum", GRAD_ACCUM), + ("eval_batch_size", EVAL_BATCH_SIZE), ("lora r / alpha / dropout", f"{LORA_R} / {LORA_ALPHA} / {LORA_DROPOUT}"), ("lora scaling (alpha/r)", LORA_ALPHA / LORA_R), ("eval_steps / save_steps", f"{EVAL_STEPS} / {SAVE_STEPS}"), @@ -228,6 +245,21 @@ def load_examples(tokenizer, path, label): return Dataset.from_dict({"input_ids": input_ids_list, "labels": labels_list}) +class PerStepMemoryCallback(TrainerCallback): + """Imprime el pico de memoria CUDA de CADA paso (no del promedio de la corrida) y lo + resetea, para poder correlacionar un pico puntual con la longitud del ejemplo que lo + causo. Se activa con PER_STEP_MEMORY_LOG=1 -- no cambia nada del comportamiento normal + de entrenamiento, es solo diagnostico. Existe porque un OOM de la Fase 6 con + GRAD_ACCUM=16 no permitia saber CUAL de los 16 micro-batches del paso lo causo.""" + + def on_step_end(self, args, state, control, **kwargs): + peak_gb = torch.cuda.max_memory_allocated() / (1024 ** 3) + reservado_gb = torch.cuda.max_memory_reserved() / (1024 ** 3) + print(f"[MEM] paso {state.global_step}: pico_asignado={peak_gb:.2f}GB " + f"pico_reservado={reservado_gb:.2f}GB", flush=True) + torch.cuda.reset_peak_memory_stats() + + class DataCollatorForCausalLMWithMasking: def __init__(self, pad_token_id): self.pad_token_id = pad_token_id @@ -306,6 +338,7 @@ def main(): output_dir=str(OUTPUT_DIR), num_train_epochs=NUM_EPOCHS, per_device_train_batch_size=1, + per_device_eval_batch_size=EVAL_BATCH_SIZE, gradient_accumulation_steps=GRAD_ACCUM, gradient_checkpointing=True, bf16=True, @@ -323,12 +356,17 @@ def main(): report_to="none", ) + callbacks = [] + if os.environ.get("PER_STEP_MEMORY_LOG", "").lower() in ("1", "true", "yes"): + callbacks.append(PerStepMemoryCallback()) + trainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, data_collator=DataCollatorForCausalLMWithMasking(tokenizer.pad_token_id), + callbacks=callbacks, ) torch.cuda.reset_peak_memory_stats()