mllam/neural-lam

Complete deterministic training by adding CUBLAS workspace config and DataLoader worker seeding

Open

#265 创建于 2026年2月25日

在 GitHub 查看
 (3 评论) (0 反应) (0 负责人)Python (276 fork)auto 404
enhancementhelp wanted

仓库指标

Star
 (282 star)
PR 合并指标
 (PR 指标待抓取)

描述

The Trainer already sets deterministic=True (train_model.py#L319), but two things are still missing for runs to be truly reproducible with the same seed.

  1. CUBLAS_WORKSPACE_CONFIG is not set Without this env var, cuBLAS matrix multiplications on GPU can still be non-deterministic even with deterministic=True. PyTorch documents this explicitly. Fix is one line at the top of main():
os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
  1. No worker_init_fn on the DataLoaders With num_workers=16, each worker gets a copy of the random state. Without a worker_init_fn, workers across runs can end up with correlated seeds. PyTorch recommends seeding each worker independently:
def _worker_init_fn(worker_id):
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)

Then pass worker_init_fn=_worker_init_fn to all three DataLoaders in WeatherDataModule.

贡献者指南