grad accum in LoRA distributed recipe

[ebsmothers]

Context

• We have it in all our other finetune recipes, might as well add it to this one too

Implemented following the approach in full_finetune_distributed.py

Run without grad accumulation

CUDA_VISIBLE_DEVICES=5,6 tune run --nproc_per_node 2 --rdzv-backend=c10d --rdzv-endpoint=localhost:20000 lora_finetune_distributed --config llama2/7B_lora checkpointer=torchtune.utils.FullModelTorchTuneCheckpointer checkpointer.checkpoint_dir=/data/users/ebs/checkpoints checkpointer.checkpoint_files=['llama2-7b-torchtune.pt'] checkpointer.output_dir=/data/users/ebs/checkpoints/new_tokenizer tokenizer.path=/data/users/ebs/checkpoints/lora-debug/tokenizer.model max_steps_per_epoch=20 metric_logger=torchtune.utils.metric_logging.WandBLogger metric_logger.project=testing
...
1|20|Loss: 1.3451744318008423:   0%|▎            

Can see from wandb that there are 20 iterations logged

Run with grad accumulation

CUDA_VISIBLE_DEVICES=5,6 tune run --nproc_per_node 2 --rdzv-backend=c10d --rdzv-endpoint=localhost:20000 lora_finetune_distributed --config llama2/7B_lora checkpointer=torchtune.utils.FullModelTorchTuneCheckpointer checkpointer.checkpoint_dir=/data/users/ebs/checkpoints checkpointer.checkpoint_files=['llama2-7b-torchtune.pt'] checkpointer.output_dir=/data/users/ebs/checkpoints/new_tokenizer tokenizer.path=/data/users/ebs/checkpoints/lora-debug/tokenizer.model max_steps_per_epoch=20 gradient_accumulation_steps=2 metric_logger=torchtune.utils.metric_logging.WandBLogger metric_logger.project=testing

...
1|40|Loss: 1.5071550607681274:   0%|▌ 

Since tqdm logs iteration number, not step number, we can see that both cases run for the expected number of iterations.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/644

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 9d6b4c6 with merge base 32d66df ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[RdoubleA]

This should be a utility at this point since it's in every recipe

[ebsmothers]

Yeah I just don't know where to put it tbh. None of the existing files feel relevant and I don't wanna create a new file just for this. Lmk if you have suggestions here

[RdoubleA]

maybe torchtune/utils/training.py or torchtune/utils/optim_utils.py? I don't see any harm in making a new file

[ebsmothers]

Tbh I'm inclined to go the other way and say that we should just do this inline. I definitely do not want to add even more indirection than we already have. And this is literally just doing a mod check

[ebsmothers]

Also @joecummings @kartikayk or @rohan-varma if you have thoughts here

[joecummings]

Do it inline.

[RdoubleA]

hm yeah, that's cleaner than I thought it would be. excellent choice

[kartikayk]

Very clean - thanks for adding this!

[kartikayk]

Since tqdm logs iteration number, not step number

How motivated are you to fix this? :) I think it should be a two line change tbh:

pbar = tqdm(desc=f"Training Epoch: {epoch+1}", total=self._dataloader)
....
# update pbar whenever we take a step
pbar.update(1)

I think this should work and then we'll no longer have the discrepancy between tqdm and our wndb logs. Can happen in a follow up as well.

[ebsmothers]

Since tqdm logs iteration number, not step number

How motivated are you to fix this? :) I think it should be a two line change tbh:

pbar = tqdm(desc=f"Training Epoch: {epoch+1}", total=self._dataloader)
....
# update pbar whenever we take a step
pbar.update(1)

I think this should work and then we'll no longer have the discrepancy between tqdm and our wndb logs. Can happen in a follow up as well.

I will probably just punt to a follow-up if it's all the same