Determine FSDP/deepspeed settings on device select.

[kallewoof]

Without this, the OS env check for accelerate will fail.

[NanoCode012]

Could you provide an example of the error that you're trying to fix and how to reproduce that?

[kallewoof]

It's not a direct error, but a bug. Perhaps the behavior is intended, in which case it should be fixed in some other way.

In either case, unless I'm mistaken, the if case in

https://github.com/OpenAccess-AI-Collective/axolotl/blob/7ee3c4cacb1c7a0d92810247f014a7394e07db80/src/axolotl/utils/config.py#L37-L41

will never be true, unless the user manually does ACCELERATE_USE_something=true python -m ....

because choose_device() is called before setup_trainer() in every case.

[NanoCode012]

I'm not so clear about that line. As I recall, the only env for that is ACCELERATE_USE_DEEPSPEED.

[winglian]

thanks @kallewoof, I made a change to your PR, mostly on naming the new function and calling it between the config validation and normalization step (where the device setup happens)

[kallewoof]

@winglian Looks good to me.

[kallewoof]

Actually, the normalization and validation should be swapped. Right now if you have e.g. a None eval_batch_size, you will get a warning about it being different from the batch size, even though the normalization part sets it to the same value if it is None. (This is unrelated to the fix in this PR, but it is something I've noticed and figured I'd mention since we're touching those lines.)

Edit: this:

eval_batch_size != micro_batch_size. This can lead to VRAM instability.

[winglian]

Actually, the normalization and validation should be swapped. Right now if you have e.g. a None eval_batch_size, you will get a warning about it being different from the batch size, even though the normalization part sets it to the same value if it is None. (This is unrelated to the fix in this PR, but it is something I've noticed and figured I'd mention since we're touching those lines.)

Edit: this:

eval_batch_size != micro_batch_size. This can lead to VRAM instability.

I'll address this in a different PR

[NanoCode012]

Actually, the normalization and validation should be swapped. Right now if you have e.g. a None eval_batch_size, you will get a warning about it being different from the batch size, even though the normalization part sets it to the same value if it is None. (This is unrelated to the fix in this PR, but it is something I've noticed and figured I'd mention since we're touching those lines.)
Edit: this:

eval_batch_size != micro_batch_size. This can lead to VRAM instability.

I'll address this in a different PR

I think I just fixed this recently. #896

[kallewoof]

Actually, the normalization and validation should be swapped. Right now if you have e.g. a None eval_batch_size, you will get a warning about it being different from the batch size, even though the normalization part sets it to the same value if it is None. (This is unrelated to the fix in this PR, but it is something I've noticed and figured I'd mention since we're touching those lines.)
Edit: this:

eval_batch_size != micro_batch_size. This can lead to VRAM instability.

I'll address this in a different PR

I think I just fixed this recently. #896

No, you circumvented one of the symptoms, but the underlying issue still remains. Normalization should happen before validation, to avoid having to put in code like in #896.

[winglian]

No, you circumvented one of the symptoms, but the underlying issue still remains. Normalization should happen before validation, to avoid having to put in code like in #896.

The reason we validate before normalization is because any errors that need to be surfaced to the user could be harder to understand since what they input could be indirectly different from the validation after normalization. I'm open to figuring out a better way.

[kallewoof]

That makes sense. I think the opposite is the case too, though, where a user may become unnecessarily confused because an error or warning is raised that ultimately does not reflect the actual parameters used in training. Such as the VRAM warning mentioned above.