How does Zero Stage 3 free up weight memory? Won't the computation graph reference it?

[Laiyi97]

# param.data doesn't store anything meaningful in partitioned state
param.data = torch.empty(0, dtype=param.dtype, device=param.device)

I noticed that Zero Stage 3 performs this operation, but I'm very confused—won't PyTorch's computation graph still retain references to the full weights?

[tjruwase]

@Laiyi97, assuming I understand your question correctly, computation graph would have references to the tensor container but not to tensor.data payload which holds the weight values.

[Laiyi97]

Thank you for your explanation.