The main process hangs every time

[ShaneTian]

This is the main code about the issue:

trainer.fit(model, datamodule=model.datamodule)

if hasattr(model.datamodule, "test_dataloader"):
    test_outputs = trainer.test(datamodule=model.datamodule, ckpt_path="best")
    if trainer.is_global_zero:
        test_output_file = os.path.join(trainer.log_dir, "test_outputs.json")
        with open(test_output_file, "w") as fout:
            json.dump(test_outputs, fout, indent=2)
            logger.info(f"save test outputs to `{test_output_file}`.")

if hasattr(model.datamodule, "predict_dataloader"):
    predict_outputs = trainer.predict(datamodule=model.datamodule, ckpt_path="best")
    if trainer.is_global_zero:
        predict_outputs = [item for batchs in predict_outputs for item in batchs]
        predict_output_file = os.path.join(trainer.log_dir, "predict_outputs.txt")
        with open(predict_output_file, "w") as fout:
            fout.write("\n".join(predict_outputs))
            logger.info(f"save prediction outputs to `{predict_output_file}`.")

• version: 1.6.0
• strategy: DeepSpeed ZeRO stage 2
• GPU: 8 A100

But every time, the process hangs here:

2022-04-01 17:43:49,291 INFO __main__ train.py[150]: save test outputs to `output/tmp/T5/version_2/test_outputs.json`.
initializing deepspeed distributed: GLOBAL_RANK: 0, MEMBER: 1/8

I found that the GPU memory of 0/7th is 300M less than other cards.

What happened?

[ShaneTian]

If I remove if trainer.is_global_zero, I got one progress bar

Testing DataLoader 0: : 63it [05:16,  5.03s/it]

but eight log

save test outputs to `output/tmp/T5/version_2/test_outputs.json`.
save test outputs to `output/tmp/T5/version_2/test_outputs.json`.
...

So, I do not ensure that test loop is triggered in main process or all processes

Same as predict loop.