Skip to content

Commit

Permalink
edge-casing for multiGPU hf to sequential case
Browse files Browse the repository at this point in the history
  • Loading branch information
haileyschoelkopf committed Oct 30, 2023
1 parent b02d989 commit fc17ada
Showing 1 changed file with 4 additions and 4 deletions.
8 changes: 4 additions & 4 deletions tools/ckpts/convert_hf_to_sequential.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@
"--output-dir checkpoints/neox_converted/pythia/70m",
"--cache-dir checkpoints/HF",
"--config configs/pythia/70M.yml configs/local_setup.yml",
"--test",
"--test",
]
)

Expand Down Expand Up @@ -519,15 +519,15 @@ def get_non_existing_dir(tmp_dir):
model, optimizer, _, lr_scheduler = deepspeed.initialize(
model=model,
optimizer=optimizer,
args=neox_args,
# args=neox_args,
lr_scheduler=lr_scheduler,
dist_init_required=False,
model_parameters=None,
config_params=neox_args.deepspeed_config,
mpu=mpu if not neox_args.is_pipe_parallel else None,
)

if os.environ["OMPI_COMM_WORLD_RANK"] == "0":
if os.environ.get("OMPI_COMM_WORLD_RANK", "1") == "0":
os.makedirs(f"{tmp_cache_dir}", exist_ok=True)

torch.distributed.barrier()
Expand Down Expand Up @@ -566,7 +566,7 @@ def get_non_existing_dir(tmp_dir):
print("==========================================")
convert(hf_model, ckpt_dir=ckpt_dir, output_dir=args.output_dir)

if os.environ["OMPI_COMM_WORLD_RANK"] == "0":
if os.environ.get("OMPI_COMM_WORLD_RANK", "1") == "0":
# cleanup temp dir
os.system(f"rm -r {tmp_cache_dir}")

Expand Down

0 comments on commit fc17ada

Please sign in to comment.