Skip to content

the all_reduce thing in parallel.py was having errors with different buffer lengths; this should fix this #47

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
a4894z wants to merge 1 commit into
base: main
Choose a base branch
from
Open

the all_reduce thing in parallel.py was having errors with different buffer lengths; this should fix this #47

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
13 changes: 11 additions & 2 deletions src/ptychi/parallel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,11 +104,20 @@ def sync_buffer(
buffer = buffer.unsqueeze(0)
unsqueezed = True

slicer = slice(None) if indices is None else indices
# For all_reduce, we need consistent shapes across all ranks
# so we must operate on the full buffer, not sliced views
if source_rank is not None:
slicer = slice(None) if indices is None else indices
dist.broadcast(buffer[slicer], src=source_rank)
else:
dist.all_reduce(buffer[slicer], op=op)
if indices is not None:
# Create a temporary buffer for the indexed elements
temp_buffer = torch.zeros_like(buffer)
temp_buffer[indices] = buffer[indices]
dist.all_reduce(temp_buffer, op=op)
buffer[indices] = temp_buffer[indices]
Comment on lines +114 to +118
Copy link

Copilot AI Oct 18, 2025

Choose a reason for hiding this comment

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

Creating a temporary buffer the size of the full buffer for partial updates is inefficient for memory usage. Consider implementing a gather-reduce-scatter pattern or using all_gather followed by local reduction to minimize memory overhead, especially for large buffers with small index subsets.

Suggested change
# Create a temporary buffer for the indexed elements
temp_buffer = torch.zeros_like(buffer)
temp_buffer[indices] = buffer[indices]
dist.all_reduce(temp_buffer, op=op)
buffer[indices] = temp_buffer[indices]
# Efficient gather-reduce-scatter for partial index updates
local_part = buffer[indices].clone()
gathered_parts = [torch.zeros_like(local_part) for _ in range(dist.get_world_size())]
dist.all_gather(gathered_parts, local_part)
# Perform reduction locally
stacked = torch.stack(gathered_parts, dim=0)
if op == dist.ReduceOp.SUM:
reduced = stacked.sum(dim=0)
elif op == dist.ReduceOp.PRODUCT:
reduced = stacked.prod(dim=0)
elif op == dist.ReduceOp.MIN:
reduced, _ = stacked.min(dim=0)
elif op == dist.ReduceOp.MAX:
reduced, _ = stacked.max(dim=0)
else:
raise NotImplementedError(f"Unsupported reduction op: {op}")
buffer[indices] = reduced

Copilot uses AI. Check for mistakes.
else:
dist.all_reduce(buffer, op=op)

if unsqueezed:
buffer = buffer.squeeze(0)
Expand Down