Add main grad before fwd pass

[vedanuj]

Adds main_grad before FWD pass to FlatParameter

to be used with https://github.com/fairinternal/xlformers/pull/1418

[mkardas]

Perhaps some check is needed to make sure parameters are not shared (as would be the case with weights tying)?

[mkardas]

Doesn't .main_grad need to be restored somewhere before next forward?

[awgu]

If we construct flat_param.main_grad before forward and set individual param.main_grad as view into flat_param.main_grad before forward, then we hold the flat_param.main_grad in memory for all FSDP instances going into backward, which may increase peak memory.

An alternative option is to construct flat_param.main_grad (as zeros or empty depending on if TE adds or copies to the memory) in the pre-backward hook and separately set param.main_grad as views into flat_param.main_grad only in the pre-backward hook.

fairscale/fairscale/nn/data_parallel/fully_sharded_data_parallel.py

Line 1521 in 0f2229e

def _pre_backward_hook(*unused: Any) -> None:

One option could be to add this logic to _prep_grads_for_backward():

fairscale/fairscale/nn/data_parallel/fully_sharded_data_parallel.py

Line 1561 in 0f2229e

self._prep_grads_for_backward()

[awgu]

nit: torch can upcast and copy in one kernel:

Suggested change

	param.main_grad.copy_(param.grad.float())
	param.main_grad.copy_(param.grad)

Existing

Upcast kernel + copy kernel

New

Only upcast kernel

Correctness Example

>>> t_fp32 = torch.empty((4,))
>>> t_bf16 = torch.randn((4,), dtype=torch.bfloat16)
>>> t_fp32
tensor([-8.3762e-20,  3.0801e-41, -1.3043e-16,  3.0801e-41])
>>> t_bf16
tensor([-1.3516, -0.5156, -0.6055,  0.3535], dtype=torch.bfloat16)
>>> t_fp32.copy_(t_bf16)
tensor([-1.3516, -0.5156, -0.6055,  0.3535])
>>> t_fp32
tensor([-1.3516, -0.5156, -0.6055,  0.3535])

[vedanuj]

@awgu It seems from my testing that the changes are still necessary in FlattenParamsWrapper otherwise it complains that .main_grad is not there for parameter.

@jspark1105 I have borrowed some changes from your PR #1136 to update the view when reallocating the zero buffers for main_grad.

[awgu]

Are we concerned that this param.main_grad.eq(0).all() might be a CPU sync? Perhaps, it is not so much a concern if we already have CPU syncs for rate limiting FSDP.

[vedanuj]

Is there another way I can check if main_grad is non zero without doing a CPU sync?

[jspark1105]

We are checking if this is all zeros to skip modules that didn't use main_grad?

[vedanuj]

yes .. because all parameters have .main_grad, so not sure how to make sure we are not using the ones that do not have the grads stored in .main_grad