Emil_Shanaty/micro-scout
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: track mixed precision updates and preserve CUDA resume state

Browse Source
This commit is contained in:
emil28092005
2026-09-16 04:10:40 +03:00
parent 1ac263b678
commit 9c41561977
1 changed files with 21 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/micro_scout/train.py
+21 -6
View File
@@ -112,9 +112,11 @@ def train(data: Path, output: Path, config: dict, device: str, resume: Path | No
scheduler = get_linear_schedule_with_warmup( scheduler = get_linear_schedule_with_warmup(
optimizer, int(total_steps * config["warmup_ratio"]), total_steps optimizer, int(total_steps * config["warmup_ratio"]), total_steps
) )
amp = device.startswith("cuda") and config.get("mixed_precision", True) use_cuda = encoder.device.type == "cuda"
amp = use_cuda and config.get("mixed_precision", True)
scaler = torch.amp.GradScaler("cuda", enabled=amp) scaler = torch.amp.GradScaler("cuda", enabled=amp)
start_epoch, start_batch, step, best = 0, 0, 0, -1.0 start_epoch, start_batch, step, best = 0, 0, 0, -1.0
optimizer_steps, skipped_optimizer_steps = 0, 0
if resume: if resume:
state = torch.load(resume / "training_state.pt", map_location="cpu", weights_only=True) state = torch.load(resume / "training_state.pt", map_location="cpu", weights_only=True)
if ( if (
@@ -133,7 +135,9 @@ def train(data: Path, output: Path, config: dict, device: str, resume: Path | No
state["best"], state["best"],
) )
torch.set_rng_state(state["torch_rng"]) torch.set_rng_state(state["torch_rng"])
if amp: optimizer_steps = state["optimizer_steps"]
skipped_optimizer_steps = state["skipped_optimizer_steps"]
if use_cuda:
torch.cuda.set_rng_state_all(state["cuda_rng"]) torch.cuda.set_rng_state_all(state["cuda_rng"])
atomic_json(output / "run.json", run_info) atomic_json(output / "run.json", run_info)
stop_requested = False stop_requested = False
@@ -166,6 +170,8 @@ def train(data: Path, output: Path, config: dict, device: str, resume: Path | No
"epoch": epoch, "epoch": epoch,
"batch": batch, "batch": batch,
"step": step, "step": step,
"optimizer_steps": optimizer_steps,
"skipped_optimizer_steps": skipped_optimizer_steps,
"best": best, "best": best,
"config": config, "config": config,
"dataset_manifest_sha256": manifest_hash, "dataset_manifest_sha256": manifest_hash,
@@ -174,7 +180,7 @@ def train(data: Path, output: Path, config: dict, device: str, resume: Path | No
"scheduler": scheduler.state_dict(), "scheduler": scheduler.state_dict(),
"scaler": scaler.state_dict(), "scaler": scaler.state_dict(),
"torch_rng": torch.get_rng_state(), "torch_rng": torch.get_rng_state(),
"cuda_rng": torch.cuda.get_rng_state_all() if amp else [], "cuda_rng": torch.cuda.get_rng_state_all() if use_cuda else [],
} }
torch.save(state, path / "training_state.pt.tmp") torch.save(state, path / "training_state.pt.tmp")
os.replace(path / "training_state.pt.tmp", path / "training_state.pt") os.replace(path / "training_state.pt.tmp", path / "training_state.pt")
@@ -221,6 +227,9 @@ def train(data: Path, output: Path, config: dict, device: str, resume: Path | No
scaler.update() scaler.update()
if scaler.get_scale() >= scale_before: if scaler.get_scale() >= scale_before:
scheduler.step() scheduler.step()
optimizer_steps += 1
else:
skipped_optimizer_steps += 1
step += 1 step += 1
running_loss += loss.item() running_loss += loss.item()
running_steps += 1 running_steps += 1
@@ -231,7 +240,11 @@ def train(data: Path, output: Path, config: dict, device: str, resume: Path | No
"epoch": epoch + 1, "epoch": epoch + 1,
"loss": running_loss / running_steps, "loss": running_loss / running_steps,
"learning_rate": scheduler.get_last_lr()[0], "learning_rate": scheduler.get_last_lr()[0],
"peak_vram_mb": torch.cuda.max_memory_allocated() / 2**20 if amp else 0, "peak_vram_mb": (
torch.cuda.max_memory_allocated() / 2**20 if use_cuda else 0
),
"optimizer_steps": optimizer_steps,
"skipped_optimizer_steps": skipped_optimizer_steps,
} }
) )
running_loss, running_steps = 0.0, 0 running_loss, running_steps = 0.0, 0
@@ -253,11 +266,13 @@ def train(data: Path, output: Path, config: dict, device: str, resume: Path | No
break break
result = { result = {
**run_info, **run_info,
"optimizer_steps": step, "batch_steps": step,
"optimizer_steps": optimizer_steps,
"skipped_optimizer_steps": skipped_optimizer_steps,
"best_validation_mrr": best, "best_validation_mrr": best,
"elapsed_seconds": time.monotonic() - started, "elapsed_seconds": time.monotonic() - started,
"stopped_early": stop_requested, "stopped_early": stop_requested,
"peak_vram_mb": torch.cuda.max_memory_allocated() / 2**20 if amp else 0, "peak_vram_mb": torch.cuda.max_memory_allocated() / 2**20 if use_cuda else 0,
} }
atomic_json(output / "result.json", result) atomic_json(output / "result.json", result)
record({"event": "complete", "best_validation_mrr": best}) record({"event": "complete", "best_validation_mrr": best})