[torchtrain] add gradient scaler

[tianyu-l]

Stack from ghstack (oldest at bottom):

• [torchtrain] add gradient clipping #28
• -> [torchtrain] add gradient scaler #25

As titled. Only enable the gradient scaler if mixed precision training is used with parameter dtype being fp16.

[lessw2020]

lgtm but one side issue.
We only want the grad scaler to be active if the param dtype is fp16 - otherwise we are checking for infs/nans that won't exist and adding overhead.
Cleanest imo is check that and then create a no op scaler via the enabled flag in ShardedGradScaler init set to False, or uglier, if/else the scaler code.
We are currently setting the dtype in the fsdp config, but it's not exposed to the user yet, but you could add a quick dtype check:
https://github.com/lessw2020/torchtrain/blob/2c7e48bbab123eefea2a6819320d5d2f81cec438/torchtrain/parallelisms/parallelize_llama.py#L57

[tianyu-l]

lgtm but one side issue. We only want the grad scaler to be active if the param dtype is fp16 - otherwise we are checking for infs/nans that won't exist and adding overhead. Cleanest imo is check that and then create a no op scaler via the enabled flag in ShardedGradScaler init set to False, or uglier, if/else the scaler code. We are currently setting the dtype in the fsdp config, but it's not exposed to the user yet, but you could add a quick dtype check: https://github.com/lessw2020/torchtrain/blob/2c7e48bbab123eefea2a6819320d5d2f81cec438/torchtrain/parallelisms/parallelize_llama.py#L57

Thanks for the note! I'm new to gradient scaling, and not sure if grad scaler should be used only when the param dtype is fp16.

According to the GradScaler doc (https://github.com/pytorch/pytorch/blob/main/torch/amp/grad_scaler.py#L90), a larger scale should be used to minimize gradient underflow, and only for fp16 this could cause gradient overflow. This seems to implicitly suggest that gradient scaling should be used in general and only for fp16 the Inf/NaN check is needed.

I might interpret it wrong, so it would great if you can share some paper/notes on grad scaling!

[lessw2020]

the larger scaling factor they reference though is for fp16 - you can both underflow and overflow with FP16. Amp is inherently using FP16, they just don't expressly mention it in that paragraph.

For BF16 and FP32 have the same dynamic range (only fp16 has the limited range issue) so no reason to scale/use ShardedGradeScaler for those dtypes.
I did double check with Andrew to make sure (though we would have been telling people the incorrect advise for 2 years if not correct lol) and he concurred that there is no need to use ShardedGradScaler except for fp16 case.

[lessw2020]

There's a good overview on fp16 and scaling here:
https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html#hpformat

[tianyu-l]

There's a good overview on fp16 and scaling here: https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html#hpformat

Just realized I read the paper Mixed Precision Training long ago but forgot almost everything. Good to learn it again :)

[lessw2020]

lgtm! Thanks for updating to check for fp16.
Let's see if Wanchao has any feedback as well.

[wanchaol]

lgtm!

[wanchaol]

we can probably later control the param_dtype as a TrainOption so that we don't rely this model be a FSDP instance. right now it looks good

[wanchaol]

Looks like you are using ghstack. Please note that for ghstack, you need to use ghstack land in cmdline instead of clicking squash and merge button, the latter will not merge the changes to main and to the base branch instead