Set teacher ckptr output_dir to match student in KD configs

[ebsmothers]

There are still a few cases left that will raise errors after #2181. This PR changes the rest of them (just the teacher checkpointer in our KD recipes). Since we don't write the teacher checkpoints anyways this is a no-op.

Python script to find all such instances where checkpointer.checkpoint_dir == checkpointer.output_dir:

import os
from omegaconf import OmegaConf

def main():
    for root, dirs, files in os.walk("/data/users/ebs/ebs-torchtune-alt/recipes"):
        for file in files:
            if file.endswith(".yaml"):
                # print(root, dirs, file)
                cfg = OmegaConf.load(os.path.join(root, file))
                OmegaConf.resolve(cfg)
                for k in cfg:
                    if "checkpointer" in k:
                        checkpointer_cfg = cfg.get(k, {})
                        output_dir = checkpointer_cfg.get('output_dir', None)
                        checkpoint_dir = checkpointer_cfg.get('checkpoint_dir', None)
                        if output_dir == checkpoint_dir:
                            print(os.path.join(root, file), k)

if __name__ == "__main__":
    main()

Test plan

tune run knowledge_distillation_single_device --config llama3_2/8B_to_1B_KD_lora_single_device \
max_steps_per_epoch=5 epochs=2
...
rm -r /tmp/torchtune/llama3_2_8B_to_1B/KD_lora_single_device/epoch_1
tune run knowledge_distillation_single_device --config llama3_2/8B_to_1B_KD_lora_single_device \
max_steps_per_epoch=5 epochs=2 resume_from_checkpoint=True
...
INFO:torchtune.utils._logging:Loss is initialized.
INFO:torchtune.utils._logging:Dataset and Sampler are initialized.
INFO:torchtune.utils._logging:Learning rate scheduler is initialized.
WARNING:torchtune.utils._logging: Profiling disabled.
INFO:torchtune.utils._logging: Profiler config after instantiation: {'enabled': False}
2|10|Loss: 1.7557275295257568: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:11<00:00,  2.22s/it]INFO:torchtune.utils._logging:Model checkpoint of size 2.30 GiB saved to /tmp/torchtune/llama3_2_8B_to_1B/KD_lora_single_device/epoch_1/ft-model-00001-of-00001.safetensors
INFO:torchtune.utils._logging:Adapter checkpoint of size 0.08 GiB saved to /tmp/torchtune/llama3_2_8B_to_1B/KD_lora_single_device/epoch_1/adapter_model.pt
INFO:torchtune.utils._logging:Adapter checkpoint of size 0.08 GiB saved to /tmp/torchtune/llama3_2_8B_to_1B/KD_lora_single_device/epoch_1/adapter_model.safetensors
INFO:torchtune.utils._logging:Adapter checkpoint of size 0.00 GiB saved to /tmp/torchtune/llama3_2_8B_to_1B/KD_lora_single_device/epoch_1/adapter_config.json

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 1858150 with merge base cdf5ea2 ():
💚 Looks good so far! There are no failures yet. 💚

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

[felipemello1]

it would be much better if we could pass None to all of these configs that dont need an output_dir, otherwise it gives an impression that it does something