Repository: hiyouga/LLaMA-Factory
File: src/llamafactory/cli.py (Python)

Code before:
    }

    command = sys.argv.pop(1) if len(sys.argv) > 1 else "help"
    if command == "train" and (is_env_enabled("FORCE_TORCHRUN") or (get_device_count() > 1 and not use_ray())):
        # launch distributed training
        max_restarts = os.getenv("MAX_RESTARTS", "0")
        rdzv_id = os.getenv("RDZV_ID")
        nnodes = os.getenv("NNODES", "1")
        min_nnodes = os.getenv("MIN_NNODES")
        max_nnodes = os.getenv("MAX_NNODES")
        node_rank = os.getenv("NODE_RANK", "0")
        nproc_per_node = os.getenv("NPROC_PER_NODE", str(get_device_count()))
        master_addr = os.getenv("MASTER_ADDR", "127.0.0.1")
        master_port = os.getenv("MASTER_PORT", str(find_available_port()))
        logger.info_rank0(f"Initializing {nproc_per_node} distributed tasks at: {master_addr}:{master_port}")
        if int(nnodes) > 1:
            print(f"Multi-node training enabled: num nodes: {nnodes}, node rank: {node_rank}")

        env = deepcopy(os.environ)
        if is_env_enabled("OPTIM_TORCH", "1"):
            # optimize DDP, see https://zhuanlan.zhihu.com/p/671834539
            env["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
            env["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"

        # check can_resume_from_checkpoint
        model_args, data_args, training_args, finetuning_args, generating_args = get_train_args()
        if finetuning_args.stage in [
            "rm",
            "ppo",
        ] and finetuning_args.finetuning_type in ["full", "freeze"]:
            can_resume_from_checkpoint = False
            if training_args.resume_from_checkpoint is not None:
                logger.warning_rank0("Cannot resume from checkpoint in current stage.")
                training_args.resume_from_checkpoint = None
        else:
            can_resume_from_checkpoint = True
        if rdzv_id is not None and can_resume_from_checkpoint:
            # launch elastic job with fault tolerant support when possible
            # see also https://docs.pytorch.org/docs/stable/elastic/train_script.html
            rdzv_nnodes = nnodes
            # elastic number of nodes if MIN_NNODES and MAX_NNODES are set
            if min_nnodes is not None and max_nnodes is not None:
                rdzv_nnodes = f"{min_nnodes}:{max_nnodes}"
            cmd = f"torchrun --nnodes {rdzv_nnodes} --nproc-per-node {nproc_per_node} --rdzv-id {rdzv_id} --rdzv-backend c10d --rdzv-endpoint {master_addr}:{master_port} --max-restarts {max_restarts} {launcher.__file__} {' '.join(sys.argv[1:])}".split()
            process = subprocess.run(cmd, env=env, check=True)
        else:
            # NOTE: DO NOT USE shell=True to avoid security risk
            process = subprocess.run(
                (
                    "torchrun --nnodes {nnodes} --node_rank {node_rank} --nproc_per_node {nproc_per_node} "
                    "--master_addr {master_addr} --master_port {master_port} {file_name} {args}"

Review comment (suggestion):
Thanks for this new feature! Can we remove the `can_resume_from_checkpoint` flag? Because in most cases, we can resume from the checkpoint.

Code after:
    }

    command = sys.argv.pop(1) if len(sys.argv) > 1 else "help"
    if command == "train" and (is_env_enabled("FORCE_TORCHRUN") or (get_device_count() > 1 and not use_ray())):
        # launch distributed training
        max_restarts = os.getenv("MAX_RESTARTS", "0")
        rdzv_id = os.getenv("RDZV_ID")
        nnodes = os.getenv("NNODES", "1")
        min_nnodes = os.getenv("MIN_NNODES")
        max_nnodes = os.getenv("MAX_NNODES")
        node_rank = os.getenv("NODE_RANK", "0")
        nproc_per_node = os.getenv("NPROC_PER_NODE", str(get_device_count()))
        master_addr = os.getenv("MASTER_ADDR", "127.0.0.1")
        master_port = os.getenv("MASTER_PORT", str(find_available_port()))
        logger.info_rank0(f"Initializing {nproc_per_node} distributed tasks at: {master_addr}:{master_port}")
        if int(nnodes) > 1:
            print(f"Multi-node training enabled: num nodes: {nnodes}, node rank: {node_rank}")

        env = deepcopy(os.environ)
        if is_env_enabled("OPTIM_TORCH", "1"):
            # optimize DDP, see https://zhuanlan.zhihu.com/p/671834539
            env["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
            env["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"

        if rdzv_id is not None:
            # launch elastic job with fault tolerant support when possible
            # see also https://docs.pytorch.org/docs/stable/elastic/train_script.html
            rdzv_nnodes = nnodes
            # elastic number of nodes if MIN_NNODES and MAX_NNODES are set
            if min_nnodes is not None and max_nnodes is not None:
                rdzv_nnodes = f"{min_nnodes}:{max_nnodes}"
            cmd = [
                "torchrun",
                "--nnodes",
                rdzv_nnodes,
                "--nproc-per-node",
                nproc_per_node,
                "--rdzv-id",
                rdzv_id,
                "--rdzv-backend",
                "c10d",
                "--rdzv-endpoint",
                f"{master_addr}:{master_port}",
                "--max-restarts",
                max_restarts,
                launcher.__file__,
                *sys.argv[1:],
            ]
            process = subprocess.run(cmd, env=env, check=True)
        else:
            # NOTE: DO NOT USE shell=True to avoid security risk