Implement gradient clipping

[tengyifei]

Fixes #90

The Hugging Face trainer clips gradient by norm by default. It turns out this makes a big difference in training stability and simply adding gradient clipping will bring torchprime into parity with Hugging Face.

Experiment: http://tb/share/yuBUPRF6KqZKPhwJkXNvc

Tested:

tp run --name tp-linear-clip-norm-2 torchprime/torch_xla_models/train.py model=llama-3-8b ici_mesh.fsdp=256 profile_step=3 profile_duration=30000 task.max_steps=1000 logging_steps=1 task.global_batch_size=256 dataset.hf_dataset_config_name=wikitext-103-raw-v1 run_name=tp-linear-clip-norm-2

tp run --use-hf torchprime/hf_models/train.py train_script.args.per_device_train_batch_size=256 +train_script.args.log_loss=true train_script.args.logging_strategy=steps +train_script.args.logging_steps=1 +train_script.args.logging_first_step=true +train_script.args.report_to=tensorboard train_script.args.max_steps=1000

[yaoshiang]

grad clipping by norm is not the only way to clip grads. Clipping by absolute value is valid. Can you update this PR to enable both? Also, the grad clipping should be configuring outside the optimizer. The grad clipping is upstream of the optimizer - the optimizer has no knowledge of grad clipping - can you move the grad clipping config to a top-level section?

[tengyifei]

@yaoshiang done. ptal

[yaoshiang]

looking good! I realize you are pushing hard to finish this off, if you have cycles, please add a unit test of some kind. Numerical would be ideal but might take a bit of time. You could do like a single layer, with fixed numbers, have the loss function be MAE, so now you know the exact grads, try your clip, and see if you get what you expect. If you feed those instructions to an LLM it'll probably do a decent job. But consider it optional since you have numerical proof against the HF curves that the norm version at least is working./ Thanks!