Skip to content

Fix: Implement get_device_stats() for train_controller #762

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
garrett4wade merged 1 commit into from
Dec 24, 2025
Merged

Fix: Implement get_device_stats() for train_controller #762

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions areal/controller/train_controller.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -487,6 +487,9 @@ def connect_engine(self, rollout: RolloutController, meta: WeightUpdateMeta):
self._init_weight_update_from_distributed(meta)
self.weight_update_group_initialized = True

def get_device_stats(self):
return self._custom_function_call("get_device_stats")
Comment on lines +490 to +491
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Dec 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The current implementation using _custom_function_call will only return the device statistics from the first data-parallel head worker. This is because _custom_function_call is designed for scenarios where results from workers are either identical or can be merged like tensors. For device statistics, where each worker has unique information, this approach leads to loss of data from other workers.

In a multi-worker environment, this would result in incomplete and potentially misleading statistics. A better approach is to explicitly gather the stats from all workers and return them as a list. This ensures that the caller receives a complete picture of the device status across the entire training setup.

Suggested change
def get_device_stats(self):
return self._custom_function_call("get_device_stats")
def get_device_stats(self):
"""Gets device statistics from all managed workers.
Returns:
list: A list of device statistics objects, one for each worker.
"""
tasks = [
self.scheduler.async_call_engine(worker.id, "get_device_stats")
for worker in self.workers
]
return self._run_async_task(asyncio.gather(*tasks))


def prepare_batch(
self,
dataloader: StatefulDataLoader,
Expand Down