| 12345678910111213141516171819202122232425262728 |
- import torch
- model_params_rank0 = torch.load("rank0model.pth",
- map_location = lambda storage, loc: storage.cuda(0))
- model_params_rank1 = torch.load("rank1model.pth",
- map_location = lambda storage, loc: storage.cuda(0))
- master_params_rank0 = torch.load("rank0master.pth",
- map_location = lambda storage, loc: storage.cuda(0))
- master_params_rank1 = torch.load("rank1master.pth",
- map_location = lambda storage, loc: storage.cuda(0))
- for model_rank0, model_rank1, master_rank0, master_rank1 in zip(
- model_params_rank0,
- model_params_rank1,
- master_params_rank0,
- master_params_rank1):
- assert torch.allclose(model_rank0, model_rank1), "Model param mismatch"
- assert torch.allclose(master_rank0, master_rank1), "Master param mismatch"
- # Some debugging/investigation assistance code:
- # maxval, maxind = torch.max(((torch.abs(model_rank0).float())/torch.abs(master_rank0)).view(-1), 0)
- # offending_val_half = model_rank0.view(-1)[maxind.item()]
- # offending_val_float = master_rank0.view(-1)[maxind.item()]
- # print(maxval.item(), maxind.item(), offending_val_half.item(), offending_val_float.item(),
- # offending_val_float.half().item())
- # rtol needs to be > 2^-11 because of denormals...
- assert torch.allclose(model_rank0, master_rank0.half(), rtol=.005), "Model-master mismatch"
- print("OK: Model and master params match across ranks.")
|
