fdsp config dict fix, todo list, add torchdistx support

2023-04-30 13:32:07 -04:00
parent 9190ada23a
commit ad2b48c0fa
3 changed files with 24 additions and 3 deletions
--- a/src/axolotl/utils/models.py
+++ b/src/axolotl/utils/models.py
@@ -179,6 +179,11 @@ def load_model(
                m.scales = m.scales.half()
                m.bias = m.bias.half()

+    if torch.cuda.device_count() > 1 and int(os.getenv("WORLD_SIZE", "1")) > 1:
+        model.is_parallelizable = True
+        model.model_parallel = True
+
+
    # TODO resume_from_checkpoint handling
    return model, tokenizer, lora_config