# ==========================================
# CELL 6: TRAINING ARGUMENTS
# ==========================================
# Run this AFTER cell_05

total_steps = len(dataset) // (BATCH_SIZE * GRAD_ACCUM) * EPOCHS
print(f"Estimated total steps: {total_steps:,}")
print(f"Estimated time on A4000: ~{total_steps / 2.5 / 3600:.1f} hours\n")

training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    overwrite_output_dir=True,

    num_train_epochs=EPOCHS,
    per_device_train_batch_size=BATCH_SIZE,
    gradient_accumulation_steps=GRAD_ACCUM,

    learning_rate=LR,
    lr_scheduler_type="cosine",
    warmup_steps=500,

    optim="adamw_8bit",
    weight_decay=0.01,
    max_grad_norm=0.3,

    bf16=torch.cuda.is_bf16_supported(),
    fp16=not torch.cuda.is_bf16_supported(),

    logging_steps=50,
    save_strategy="steps",
    save_steps=2000,
    save_total_limit=2,

    group_by_length=True,
    remove_unused_columns=False,

    report_to="none",  # Change to "wandb" if you want logging
    seed=3407,
)

print("Training arguments configured.")
print(f"  Effective batch size: {BATCH_SIZE * GRAD_ACCUM}")
print(f"  Learning rate: {LR}")
print(f"  Context length: {MAX_SEQ_LENGTH}")
print(f"  Epochs: {EPOCHS}")