foundation-model-stack / fms-hf-tuning

🚀 Collection of tuning recipes with HuggingFace SFTTrainer and PyTorch FSDP.
Apache License 2.0
28 stars 48 forks source link

fix: add enable_aim build args in all stages needed #337

Closed anhuong closed 2 months ago

anhuong commented 2 months ago

Description of the change

When I tried to build the image with AIM, it installed AIM correctly but didn't have the aim directory from the final stage since the build arg wasn't available in the final stage. This fixes the issue by defining enable_aim at the top and then enabling the flag in the stages where it is used.

Related issue number

How to verify the PR

Was the PR tested