Skip to content

Add W&B worker GPU system metrics#1335

Closed
EazyReal wants to merge 1 commit into
areal-project:mainfrom
VmaxAI:codex/wandb-worker-system-metrics
Closed

Add W&B worker GPU system metrics#1335
EazyReal wants to merge 1 commit into
areal-project:mainfrom
VmaxAI:codex/wandb-worker-system-metrics

Conversation

@EazyReal

Copy link
Copy Markdown
Contributor

Summary

  • Add stats_logger.wandb.system_metrics to opt into worker-side W&B system metrics.
  • Initialize the controller W&B run as the primary shared-mode client and attach selected worker RPC processes/Ray actors as non-primary clients.
  • Document the single-controller GPU telemetry setup and regenerate CLI reference docs.

Root Cause

In single-controller training the process that calls wandb.init() can run on a CPU/controller host while actor and rollout workers run on GPU nodes. W&B system metrics are collected from initialized W&B clients, so wandb.mode=online on only the controller records controller process memory but not worker GPU utilization.

Behavior

Users can enable GPU worker telemetry with:

stats_logger:
  wandb:
    mode: shared
    system_metrics:
      enabled: true
      roles: [actor, rollout]
      gpu_device_ids: null

Workers use x_primary=False, a role/rank x_label, and x_update_finish_state=False, so they contribute system metrics without writing primary run files or finishing the run.

Validation

  • uv run ruff check areal/api/cli_args.py areal/utils/stats_logger.py areal/utils/wandb_system_metrics.py areal/infra/rpc/rpc_server.py areal/infra/rpc/ray_rpc_server.py areal/infra/scheduler/local.py areal/infra/scheduler/slurm.py areal/infra/scheduler/ray.py areal/trainer/rl_trainer.py areal/trainer/sft_trainer.py areal/trainer/dpo_trainer.py areal/trainer/rw_trainer.py tests/test_wandb_system_metrics.py
  • uv run pytest tests/test_wandb_system_metrics.py tests/test_trackio_backend.py

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant