Optimize balancer and setup debug logger.#308
Open
JacoCheung wants to merge 10 commits intoNVIDIA:mainfrom
Open
Optimize balancer and setup debug logger.#308JacoCheung wants to merge 10 commits intoNVIDIA:mainfrom
JacoCheung wants to merge 10 commits intoNVIDIA:mainfrom
Conversation
Greptile OverviewGreptile SummaryThis PR optimizes the balancer overhead and batch allgather communication in the training pipeline by implementing two key improvements: 1. Background Thread for Balancer: Offloaded H2D transfer and batch shuffling (including Karmarkar-Karp load balancing) to a background thread using 2. Fused KJT AllGather: Replaced per-KJT gathering with Additional improvements:
Confidence Score: 4/5
Important Files Changed
Sequence DiagramsequenceDiagram
participant Main as Main Thread
participant BG as Background Thread
participant GPU as GPU Stream
participant NCCL as NCCL (DP Group)
Note over Main,NCCL: Batch i+2 Processing (Optimized)
Main->>Main: _next_batch(dataloader)
Main->>BG: submit(_h2d_and_shuffle, batch)
activate BG
BG->>GPU: _to_device (H2D on memcpy_stream)
BG->>NCCL: AllGather workloads
BG->>NCCL: Karmarkar-Karp partitioning
BG->>NCCL: AllGather KJTs (fused, 2 calls)
BG->>NCCL: AllGather dense tensors
deactivate BG
Note over Main: Batch i Processing (overlapped)
Main->>GPU: Forward pass (default stream)
Main->>BG: future.result() [WAIT HERE]
activate BG
BG-->>Main: shuffled_batch
deactivate BG
Note over Main,NCCL: Safe to use DP communicator
Main->>NCCL: AllReduce loss (DP group)
Main->>GPU: Backward pass
Main->>NCCL: AllReduce gradients (DP group)
|
JacoCheung
force-pushed
the
junzhang/opt_balancer
branch
from
8b68ea4 to
2e3a262
Compare
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
examples/commons/utils/logger.py
Outdated
Comment on lines
64
to
65
| def print_rank_all(message, level=logging.INFO): | ||
| """If distributed is initialized, print only on rank 0.""" |
There was a problem hiding this comment.
Docstring is incorrect - says "print only on rank 0" but print_rank_all actually prints on all ranks.
Suggested change
| def print_rank_all(message, level=logging.INFO): | |
| """If distributed is initialized, print only on rank 0.""" | |
| """If distributed is initialized, print on all ranks.""" |
examples/commons/utils/logger.py
Outdated
Comment on lines
72
to
73
| def info_rank_all(message): | ||
| """If distributed is initialized, print only on rank 0.""" |
There was a problem hiding this comment.
Docstring is incorrect - says "print only on rank 0" but info_rank_all actually prints on all ranks.
Suggested change
| def info_rank_all(message): | |
| """If distributed is initialized, print only on rank 0.""" | |
| """If distributed is initialized, print on all ranks.""" |
examples/commons/utils/logger.py
Outdated
Comment on lines
80
to
81
| def debug_rank_all(message): | ||
| """If distributed is initialized, print only on rank 0.""" |
There was a problem hiding this comment.
Docstring is incorrect - says "print only on rank 0" but debug_rank_all actually prints on all ranks.
Suggested change
| def debug_rank_all(message): | |
| """If distributed is initialized, print only on rank 0.""" | |
| """If distributed is initialized, print on all ranks.""" |
Comment on lines
+719
to
+722
| self._shuffle_executor = concurrent.futures.ThreadPoolExecutor( | ||
| max_workers=1, thread_name_prefix="shuffle" | ||
| ) | ||
| self._shuffle_future: Optional["concurrent.futures.Future[Optional[In]]"] = None |
There was a problem hiding this comment.
Consider adding cleanup for _shuffle_executor to ensure graceful thread termination:
def __del__(self):
if hasattr(self, '_shuffle_executor'):
self._shuffle_executor.shutdown(wait=True)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR aims to fully hide the balancer overhead (KK algorithm on host) and optimize the comms of batch allgather.
Besides, a set of helper utils is added.
- The peak sol is about fwd 60% , bwd %50.