tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

fix model parallel (#816)

Browse Source
This commit is contained in:
Wing Lian
2023-11-02 21:34:22 -04:00
committed by GitHub
parent 10388a8daf
commit 964d858da0
1 changed files with 1 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/axolotl/utils/models.py
View File

@@ -442,14 +442,7 @@ def load_model(
if cfg.ddp and not load_in_8bit:
model.to(f"cuda:{cfg.local_rank}")
if (
torch.cuda.device_count() > 1
and int(os.getenv("WORLD_SIZE", "1")) > 1
and (cfg.load_in_4bit)
):
# llama is PROBABLY model parallelizable, but the default isn't that it is
# so let's only set it for the 4bit, see
# https://github.com/johnsmith0031/alpaca_lora_4bit/blob/08b3fca4a4a9e0d3945be1bab4529f100a428636/finetune.py#L130-L133
if torch.cuda.device_count() > 1 and int(os.getenv("WORLD_SIZE", "1")) == 1:
setattr(model, "is_parallelizable", True)
setattr(model, "model_parallel", True)