We fine-tuned Llama 405B on AMD GPUs

Re:We fine-tuned Llama 405B on AMD GPUs by ngcc_hk on Monday September 23, 2024 @06:30PM (#41631278)

Given it is a migration, is there actual comparison of the same model on PyTorch vs your version. The comparison table there seems to be on technical side.

Also any technical issues encountered?

[ Parent | on HN ]

Re:We fine-tuned Llama 405B on AMD GPUs by felarof on Monday September 23, 2024 @10:28PM (#41632574)

We have a few technical issues that we still need to address:

1) This entire fine-tuning run was done in JAX eager mode. I kept running out of memory (OOM) when trying to `jax.jit` the entire training step. Even gradual `jax.jit` didn't work.

2) The current version doesn't have gradient accumulation, and with a batch size of just 16, that’s not ideal. I'm working on implementing gradient accumulation next.

3) We still haven't found a good way to load large sequence-length data (like 32k sequence length). Currently, before sharding the training batch across GPUs, it ends up loading the entire batch onto a single GPU’s VRAM and causes OOM issues.

[ Parent | on HN ]

Re:We fine-tuned Llama 405B on AMD GPUs by logicchains on Tuesday September 24, 2024 @03:04AM (#41633835)

> I kept running out of memory (OOM) when trying to `jax.jit` the entire training step. Even gradual `jax.jit` didn't work.

Were you using activation checkpointing? https://jax.readthedocs.io/en/latest/_autosummary/jax.checkp... is very important for keeping memory usage reasonable when training large models.

[ Parent | on HN ]