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

Code before:
            super().setup(args, state, model, **kwargs)
            try:
                if hasattr(self, "_swanlab"):
                    swanlab_public_config = self._swanlab.get_run().public.json()
                else:  # swanlab <= 0.4.9
                    swanlab_public_config = self._experiment.get_run().public.json()
            except Exception:
                swanlab_public_config = {}

            with open(os.path.join(args.output_dir, SWANLAB_CONFIG), "w") as f:
                f.write(json.dumps(swanlab_public_config, indent=2))

    swanlab_callback = SwanLabCallbackExtension(
        project=finetuning_args.swanlab_project,
        workspace=finetuning_args.swanlab_workspace,
        experiment_name=finetuning_args.swanlab_run_name,
        mode=finetuning_args.swanlab_mode,
        config={"Framework": "🦙LlamaFactory"},
        logdir=finetuning_args.swanlab_logdir,
        tags=["🦙LlamaFactory"],
    )
    return swanlab_callback


def get_placement_group(num_workers, strategy):
    bundle = {"CPU": 1}
    if is_torch_cuda_available():
        bundle["GPU"] = 1
    elif is_torch_npu_available():
        bundle["NPU"] = 1

    bundles = [bundle for _ in range(num_workers)]
    pg = placement_group(bundles, strategy=strategy)

    return pg, bundle


def get_ray_remote_config_for_worker(
    placement_group,
    bundle_idx,
    rank,
    world_size,
    master_addr,
    master_port,
):
    # "RANK": str(rank),
    env_vars = {
        "RANK": str(rank),
        "WORLD_SIZE": str(world_size),
        "MASTER_ADDR": master_addr,
        "MASTER_PORT": master_port,

Review comment (performance):
please double check cpu:1 is reasonable or not？ Sometimes it will use bundle = {"CPU": max_number}, restricting cpu to 1 may reduce performance.

Code after:
                if hasattr(self, "_swanlab"):
                    swanlab_public_config = self._swanlab.get_run().public.json()
                else:  # swanlab <= 0.4.9
                    swanlab_public_config = self._experiment.get_run().public.json()
            except Exception:
                swanlab_public_config = {}

            with open(os.path.join(args.output_dir, SWANLAB_CONFIG), "w") as f:
                f.write(json.dumps(swanlab_public_config, indent=2))

    swanlab_callback = SwanLabCallbackExtension(
        project=finetuning_args.swanlab_project,
        workspace=finetuning_args.swanlab_workspace,
        experiment_name=finetuning_args.swanlab_run_name,
        mode=finetuning_args.swanlab_mode,
        config={"Framework": "🦙LlamaFactory"},
        logdir=finetuning_args.swanlab_logdir,
        tags=["🦙LlamaFactory"],
    )
    return swanlab_callback


def get_placement_group(num_workers: int) -> tuple["PlacementGroup", dict[str, int]]:
    r"""Get the Ray placement group for distributed training."""
    bundle = {"CPU": 10}
    device_name = get_device_name().upper()
    if device_name != "CPU":
        bundle[device_name] = 1
    bundles = [bundle for _ in range(num_workers)]
    pg = placement_group(bundles, strategy="PACK")

    return pg, bundle


def get_ray_remote_config_for_worker(
    placement_group: "PlacementGroup",
    bundle_idx: int,
    rank: int,
    world_size: int,
    master_addr: str,
    master_port: str,
    env: dict[str, str] = None,
) -> dict[str, Any]:
    r"""Get the remote config for a Ray worker."""
    env_vars = {
        "RANK": str(rank),
        "WORLD_SIZE": str(world_size),
        "MASTER_ADDR": master_addr,
        "MASTER_PORT": master_port,
        "TORCHELASTIC_USE_AGENT_STORE": "False",
    }