Instruction.md:
You are provided with a PyTorch training script at `/app/train.py`. The script trains a simple neural network using gradient accumulation and Automatic Mixed Precision (AMP). It is currently configured to simulate a server crash immediately after processing micro-batch 103 out of 200, while partway through a gradient-accumulation cycle.
Currently, when the training loop crashes and resumes from `/app/checkpoint.pt`, the loss curve diverges compared to an uninterrupted run. The checkpointing logic is failing to capture the complete mathematical state required for bit-for-bit reproducibility upon resumption.
Your task is to diagnose the cause of the divergence and fix the checkpointing logic in `/app/train.py`. You must accurately capture and restore all missing state.
The final script must be able to crash, resume from `/app/checkpoint.pt`, and finish training without restarting from the beginning or replaying micro-batches that were completed before the checkpoint. The resumed execution must continue from the saved micro-batch position and restore the complete mathematical training state required for bit-for-bit reproducibility. It must produce `/app/final_loss.txt` and `/app/model_final.pt` that are mathematically identical to those produced by an uninterrupted run. The final model checkpoint must also contain sufficient execution metadata to verify that an interrupted run genuinely resumed from the saved checkpoint rather than restarting training from micro-batch 1.
train.py:
You are provided with a PyTorch training script at `/app/train.py`. The script trains a simple neural network using gradient accumulation and Automatic Mixed Precision (AMP). It is currently configured to simulate a server crash immediately after processing micro-batch 103 out of 200, while partway through a gradient-accumulation cycle.
Currently, when the training loop crashes and resumes from `/app/checkpoint.pt`, the loss curve diverges compared to an uninterrupted run. The checkpointing logic is failing to capture the complete mathematical state required for bit-for-bit reproducibility upon resumption.
Your task is to diagnose the cause of the divergence and fix the checkpointing logic in `/app/train.py`. You must accurately capture and restore all missing state.
The final script must be able to crash, resume from `/app/checkpoint.pt`, and finish training without restarting from the beginning or replaying micro-batches that were completed before the checkpoint. The resumed execution must continue from the saved micro-batch position and restore the complete mathematical training state required for bit-for-bit reproducibility. It must produce `/app/final_loss.txt` and `/app/model_final.pt` that are mathematically identical to those produced by an uninterrupted run. The final model checkpoint must also contain sufficient execution metadata to verify that an interrupted run genuinely resumed from the saved checkpoint rather than restarting training from micro-batch 1.
solve.py:
import os
import sys
import torch
from pathlib import Path
train_script_path = Path("/app/train.py")
content = train_script_path.read_text()
# Replacement complete save_checkpoint function
new_save = ''' def save_checkpoint(step):
grads = {
name: param.grad.detach().clone()
for name, param in model.named_parameters()
if param.grad is not None
}
state = {
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"scheduler": scheduler.state_dict(),
"scaler": scaler.state_dict(),
"rng_cpu": torch.get_rng_state(),
"augmentation_rng": augmentation_generator.get_state(),
"grads": grads,
"step": step,
}
torch.save(state, checkpoint_path)'''
# Replacement complete load_checkpoint function
new_load = ''' def load_checkpoint():
if not os.path.exists(checkpoint_path):
return 0
checkpoint = torch.load(
checkpoint_path,
weights_only=False
)
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])
scheduler.load_state_dict(checkpoint["scheduler"])
scaler.load_state_dict(checkpoint["scaler"])
torch.set_rng_state(checkpoint["rng_cpu"])
augmentation_generator.set_state(checkpoint["augmentation_rng"])
for name, param in model.named_parameters():
if name in checkpoint["grads"]:
param.grad = checkpoint["grads"][name].clone()
print(f"Resuming from micro-batch {checkpoint['step']}")
return checkpoint["step"]'''
# Locate and replace the old naive functions
old_save_start = content.find(" def save_checkpoint(step):")
old_save_end = content.find(" def load_checkpoint():")
old_load_start = old_save_end
old_load_end = content.find(" # Generate the fixed synthetic training dataset.")
if old_save_start != -1 and old_load_end != -1:
fixed_content = (
content[:old_save_start]
+ new_save + "\n\n"
+ new_load + "\n\n"
+ content[old_load_end:]
)
train_script_path.write_text(fixed_content)
print("Successfully patched /app/train.py with state restoration!")
else:
print("Error: Could not locate save/load functions to patch.", file=sys.stderr)
sys.exit(1)
test_output.py:
import os
import subprocess
from pathlib import Path
import pytest
import torch
def assert_dicts_equal(d1, d2, path=""):
"""Recursively compare two dictionaries containing tensors and primitives."""
assert type(d1) == type(d2), f"Type mismatch at {path}: {type(d1)} vs {type(d2)}"
if isinstance(d1, dict):
assert set(d1.keys()) == set(d2.keys()), f"Keys mismatch at {path}"
for k in d1:
assert_dicts_equal(d1[k], d2[k], path=f"{path}.{k}" if path else str(k))
elif isinstance(d1, (list, tuple)):
assert len(d1) == len(d2), f"Length mismatch at {path}"
for i, (v1, v2) in enumerate(zip(d1, d2)):
assert_dicts_equal(v1, v2, path=f"{path}[{i}]")
elif isinstance(d1, torch.Tensor):
assert torch.equal(d1, d2), f"Tensor mismatch at {path}"
else:
assert d1 == d2, f"Value mismatch at {path}: {d1} vs {d2}"
@pytest.fixture(scope="module")
def execution_results():
"""
Executes the control run and the crash/resume run once for the entire test module.
Caches the final states so we don't have to run training 9 times.
"""
script_path = "/app/train.py"
artifacts = ["/app/checkpoint.pt", "/app/model_final.pt", "/app/final_loss.txt", "/app/crashed_once.flag"]
# --- Run A: Control Run (Unbroken) ---
for f in artifacts:
if os.path.exists(f):
os.remove(f)
subprocess.run(
["python3", script_path],
env=dict(os.environ, ALLOW_CRASH="0"),
check=True,
timeout=180,
)
control_loss = Path("/app/final_loss.txt").read_text().strip()
control_state = torch.load("/app/model_final.pt", weights_only=False)
# --- Run B: Resumed Run (Crashes at micro-batch 103) ---
for f in artifacts:
if os.path.exists(f):
os.remove(f)
# First attempt: Verify intentional crash occurs
crash_result = subprocess.run(
["python3", script_path],
env=dict(os.environ, ALLOW_CRASH="1"),
timeout=180,
)
checkpoint_exists = os.path.exists("/app/checkpoint.pt")
# Second attempt: resumes from checkpoint
subprocess.run(
["python3", script_path],
env=dict(os.environ, ALLOW_CRASH="1"),
check=True,
timeout=180,
)
resumed_loss = Path("/app/final_loss.txt").read_text().strip()
resumed_state = torch.load("/app/model_final.pt", weights_only=False)
return {
"crash_returncode": crash_result.returncode,
"checkpoint_exists": checkpoint_exists,
"control_loss": control_loss,
"resumed_loss": resumed_loss,
"control_state": control_state,
"resumed_state": resumed_state,
}
# --- Atomic Tests ---
def test_intentional_crash_occurs(execution_results):
"""Verifies that the script crashes mid-training as expected when ALLOW_CRASH=1."""
assert execution_results["crash_returncode"] != 0, "The expected crash at micro-batch 103 did not occur."
def test_checkpoint_created_before_crash(execution_results):
"""Verifies the checkpoint file was successfully written before the training script crashed."""
assert execution_results["checkpoint_exists"], "Failed to create /app/checkpoint.pt before crashing."
def test_resume_starts_from_checkpoint(execution_results):
"""
Verifies that the resumed execution genuinely starts from the saved
micro-batch position rather than restarting from step 0.
"""
control_start = execution_results["control_state"]["run_start_step"]
resume_start = execution_results["resumed_state"]["run_start_step"]
assert control_start == 0, (
f"Control run should start at 0, but started at {control_start}."
)
assert resume_start == 103, (
f"Expected resume to start at micro-batch 103, "
f"but started at {resume_start}."
)
def test_final_loss_matches(execution_results):
"""Verifies exact mathematical equality of the final loss between the control and resumed runs."""
assert execution_results["control_loss"] == execution_results["resumed_loss"], "Final loss diverged."
def test_micro_batches_completed(execution_results):
"""Verifies both the baseline and resumed training runs completed exactly 200 micro-batches."""
assert execution_results["control_state"]["micro_batches_completed"] == 200, "Control run did not complete 200 micro-batches."
assert execution_results["resumed_state"]["micro_batches_completed"] == 200, "Resumed run did not complete 200 micro-batches."
def test_model_weights_match(execution_results):
"""Verifies bit-for-bit reproducibility of the model weights after an interrupted run."""
ctrl, res = execution_results["control_state"]["model"], execution_results["resumed_state"]["model"]
for k in ctrl:
assert torch.equal(ctrl[k], res[k]), f"Model weights for layer '{k}' diverged upon resumption."
def test_optimizer_state_matches(execution_results):
"""Verifies exact equality of the optimizer momentum buffers and state."""
assert_dicts_equal(execution_results["control_state"]["optimizer"], execution_results["resumed_state"]["optimizer"], path="optimizer")
def test_scheduler_state_matches(execution_results):
"""Verifies the learning rate scheduler state is perfectly restored."""
assert execution_results["control_state"]["scheduler"] == execution_results["resumed_state"]["scheduler"], "Scheduler state diverged."
def test_scaler_state_matches(execution_results):
"""Verifies the AMP GradScaler scale multiplier and state are perfectly restored."""
assert execution_results["control_state"]["scaler"] == execution_results["resumed_state"]["scaler"], "GradScaler state diverged."