gistfile0.txt

from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP

import sys
import torch
from torch.nn import init
import math

torch.set_printoptions(precision=20)

torch.manual_seed(12)
print("RAND", torch.rand(1))
if sys.argv[1] == "0":
    t = torch.rand(100, 100)
elif sys.argv[1] == "1":
    t = torch.rand(100* 100 * 100)
    t = t[:100*100].reshape(100, 100)
else:
    assert 0

init.uniform_(t)
#init.kaiming_uniform_(t, a=math.sqrt(5))
print("NORM", t.norm())

print("RAND", torch.rand(1))

cc @myleott, a very interesting behavior that affects FSDP in flattening case for weight init.