| 142 | return candidate_files[-1] |
| 143 | |
| 144 | def run_command(cmd, description): |
| 145 | # empirical max dataLoader throughput settings (I used 1-6 H100s) |
| 146 | env = os.environ.copy() |
| 147 | env.setdefault("NG_NUM_WORKERS", str(max(2, (os.cpu_count() or 4) - 2))) |
| 148 | env.setdefault("NG_PREFETCH_FACTOR", "4") |
| 149 | env.setdefault("NG_PIN_MEMORY", "1") |
| 150 | env["NG_PERSISTENT_WORKERS"] = "0" |
| 151 | env.setdefault("TORCH_CUDNN_V8_API_ENABLED", "1") |
| 152 | |
| 153 | try: |
| 154 | result = subprocess.run(cmd, check=True, capture_output=False, env=env) |
| 155 | return True |
| 156 | except subprocess.CalledProcessError as e: |
| 157 | print(f"Error: {e.stderr}") |
| 158 | return False |
| 159 | except KeyboardInterrupt: |
| 160 | return False |
| 161 | |
| 162 | def save_training_state(model, optimizer, scheduler, config, checkpoints_dir, prefix, step): |
| 163 | """Save a checkpoint with model/optimizer/scheduler and the exact config. |