Skip to content

feat: distributed gradient gathering for multi-GPU parallelism #649

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

Open
shivam-MBZUAI wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

feat: distributed gradient gathering for multi-GPU parallelism #649

Show file tree
Hide file tree
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 neurons/base_node.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,9 @@ def __init__(self):
self._prof = None
# self.config, self.world_size … come from the concrete node

# Continue MRO chain to initialize parent classes (e.g., Trainer)
super().__init__()

async def main(self):
loop = asyncio.get_running_loop()
self._setup_signal_handlers(loop)
Expand Down
138 changes: 119 additions & 19 deletions neurons/miner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,12 +28,18 @@
import sys
import time
from datetime import datetime, timedelta, timezone

# Set CUDA memory allocator config before importing torch
if "PYTORCH_CUDA_ALLOC_CONF" not in os.environ:
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

from types import SimpleNamespace
from typing import cast

import bittensor as bt
import numpy as np
import torch
import torch.distributed as dist
import uvloop
from torch.amp.grad_scaler import GradScaler
from torch.distributed.tensor import DTensor as DT
Expand Down Expand Up @@ -465,6 +471,26 @@ async def run(self):
await tplr.neurons.update_peers(
instance=self, window=step_window, peer_start=peer_start
)
# Refresh commitments to get updated bucket info for all peers
self.comms.commitments = await self.comms.get_commitments()

# Broadcast peer / reserve lists and commitments from master to all ranks for distributed gather
if dist_helper.world_size > 1 and dist_helper.is_distributed():
payload = (
[self.comms.peers, self.comms.reserve_peers, self.comms.commitments]
if self.is_master
else [None, None, None]
)
dist.broadcast_object_list(payload, src=0)
self.comms.peers = payload[0]
self.comms.reserve_peers = payload[1]
self.comms.commitments = payload[2]

if not self.is_master:
tplr.logger.info(
f"[Rank {dist_helper.rank}] Received peers from master: gather={self.comms.peers}, reserve={self.comms.reserve_peers}"
)

Comment thread
coderabbitai[bot] marked this conversation as resolved.
peer_update_time = tplr.T() - peer_start

# 2. Load data
Expand Down Expand Up @@ -525,6 +551,15 @@ async def run(self):
else:
tplr.logger.info("Start accumulating...")

# Aggressive memory cleanup before training to prevent OOM
# This ensures memory from previous outer step is fully released
torch.cuda.empty_cache()
torch.cuda.synchronize()

# Clear any lingering gradients
if hasattr(self, "model") and self.model is not None:
self.model.zero_grad(set_to_none=True)

res = await self.inner_steps(
loader=self.loader, step_window=step_window, null_round=null_round
)
Expand Down Expand Up @@ -702,30 +737,81 @@ async def run(self):
gather_time = 0.0
should_update = True

if self.is_master:
gather_start = tplr.T()
tplr.logger.info("Waiting on gather task...")
# Use distributed gather automatically when running with multiple ranks
use_distributed_gather = (
dist_helper.world_size > 1 and dist_helper.is_distributed()
)

gather_start = tplr.T()

# Log gather mode
if use_distributed_gather:
tplr.logger.info(
f"[Rank {dist_helper.rank}] Starting distributed gather from {len(self.comms.peers)} peer(s) (world_size={dist_helper.world_size})"
)
elif self.is_master:
tplr.logger.info(
f"Starting sequential gather from {len(self.comms.peers)} peer(s)"
)

# For distributed gather, all ranks must participate (for collective operations)
# For sequential gather, only master rank performs the operation
if self.is_master or use_distributed_gather:
gather_result = await self.comms.gather_with_reserve(
my_uid=self.uid,
gather_uids=self.comms.peers,
reserve_uids=self.comms.reserve_peers,
reserve_uids=self.comms.reserve_peers, # Miners gather same way as validator
return_partials=False, # Always return merged result (distributed gather merges internally)
window=step_window,
key="gradient",
timeout=90,
device=str(self.device),
local=False,
stale_retention=100,
totalks=self.totalks,
compressor=self.compressor,
time_min=time_min,
time_max=time_max,
expected_compressed_params=self.expected_compressed_params,
)
tplr.logger.info("Gather task completed!")
gather_time = tplr.T() - gather_start
should_update = gather_result is not None

# Broadcast whether we should update to all ranks
# Log gather completion
Comment thread
shivam-MBZUAI marked this conversation as resolved.
if use_distributed_gather:
if gather_result is not None:
tplr.logger.info(
f"[Rank {dist_helper.rank}] Distributed gather complete: {len(gather_result.uids)}/{len(self.comms.peers)} successful, "
f"{len(gather_result.skipped_uids)} skipped, success_rate={gather_result.success_rate:.2%}, "
f"time={gather_time:.2f}s"
)
else:
tplr.logger.warning(
f"[Rank {dist_helper.rank}] Distributed gather failed - no gradients collected from peers"
)
else:
# Sequential gather logging (master only)
if gather_result is not None:
tplr.logger.info(
f"Sequential gather complete: {len(gather_result.uids)}/{len(self.comms.peers)} successful, "
f"{len(gather_result.skipped_uids)} skipped, "
f"success_rate={gather_result.success_rate:.2%}, "
f"time={gather_time:.2f}s"
)
else:
tplr.logger.warning(
"Sequential gather failed - no gradients collected from peers"
)

# For distributed gather: only master checks result, others will follow master's decision
# For sequential gather: only master has result anyway
if use_distributed_gather:
# Master checks if gather succeeded, non-master ranks defer to master
should_update = (
gather_result is not None if self.is_master else True
)
else:
should_update = gather_result is not None

# Broadcast whether we should update to all ranks (master's decision)
should_update = dist_helper.all_ok(
should_update, self.device, "gather_update"
)
Expand All @@ -750,6 +836,11 @@ async def run(self):
self.global_step += (
1 # Increment only when we actually do an outer step
)

# Clear gradients after outer step
if hasattr(self, "model") and self.model is not None:
self.model.zero_grad(set_to_none=True)

model_update_time = tplr.T() - update_start
if gradient_fingerprint is not None:
tplr.logger.info(
Expand Down Expand Up @@ -801,12 +892,13 @@ async def run(self):

# Add successful peers information
if gather_result is not None:
successful_uids = set(gather_result.uids)
skipped_uids = set(gather_result.skipped_uids)

debug_dict["successful_peers"] = sorted(
list(set(self.comms.peers) - set(gather_result.skipped_uids))
)
debug_dict["skipped_peers"] = sorted(
list(gather_result.skipped_uids)
list(set(self.comms.peers) - skipped_uids)
)
debug_dict["skipped_peers"] = sorted(list(skipped_uids))

# Store the debug dictionary
await self.comms.put(
Expand Down Expand Up @@ -841,9 +933,14 @@ async def run(self):
sum(momentum_norms) / len(momentum_norms) if momentum_norms else 0
)
window_total_time = tplr.T() - window_start
gather_success_rate = (
gather_result.success_rate * 100 if gather_result else 0.0
)

# Calculate success rate and extract skipped UIDs
if gather_result is None:
gather_success_rate = 0.0
skipped_uids_list = []
else:
gather_success_rate = gather_result.success_rate * 100
skipped_uids_list = list(gather_result.skipped_uids)
inner_lr = self.inner_scheduler.get_last_lr()[0]

# Only log to WandB when we've performed an outer step
Expand Down Expand Up @@ -907,9 +1004,7 @@ async def run(self):
"n_gather_peers": int(len(self.comms.peers)),
"gather_success_rate": gather_success_rate,
"gather_peers": json.dumps(self.comms.peers),
"skipped_peers": json.dumps(
gather_result.skipped_uids if gather_result else []
),
"skipped_peers": json.dumps(skipped_uids_list),
"window_total_time": window_total_time,
"peer_update_time": peer_update_time,
"compression_time": compression_time,
Expand All @@ -926,10 +1021,15 @@ async def run(self):

dist_helper.safe_barrier("post_outer_step", self.local_rank)

# Delete any remaining local variables to clear up memory
# Delete local variables and force cleanup
del shard_gradient
if gather_result is not None:
# Clear state_dict to free memory
if hasattr(gather_result, "state_dict"):
if hasattr(gather_result.state_dict, "__dict__"):
gather_result.state_dict.__dict__.clear()
del gather_result
gc.collect()
torch.cuda.empty_cache()
# Check memory threshold periodically
self.check_memory_threshold(threshold_gb=0.5)
Expand Down
6 changes: 5 additions & 1 deletion neurons/trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -884,7 +884,9 @@ async def inner_steps(

# Unscale, clip, then step via GradScaler if using fp16
self.scaler.unscale_(self.inner_optimizer)
torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
total_grad_norm = torch.nn.utils.clip_grad_norm_(
self.model.parameters(), 1.0
).item()
self.scaler.step(self.inner_optimizer)
self.scaler.update()

Expand All @@ -906,6 +908,7 @@ async def inner_steps(
else:
# Spin-up: don't step optimizer/scheduler, just clear gradients
self.scaler.update()
total_grad_norm = 0.0

self.inner_optimizer.zero_grad(set_to_none=True)

Expand All @@ -922,6 +925,7 @@ async def inner_steps(
tplr.logger.info(
f"Inner Step {inner_step_count}, "
f"Batch {batch_count}, loss: {log_loss:.4f}, "
f"grad_norm: {total_grad_norm:.4f}, "
f"accum: {accum_batch_size}/{self.hparams.batch_size}"
)
if window_entry_loss == 0.0:
Expand Down
Loading