Adding prefill decode markers to trace and enable shapes #156
Conversation
There was a problem hiding this comment.
Pull request overview
This pull request adds profiler markers to distinguish between prefill and decode phases in PyTorch profiler traces, and introduces a new environment variable to enable shape recording independently of other detailed profiling features.
Changes:
- Added
ATOM_PROFILER_SHAPESenvironment variable to control shape recording in the profiler - Added nested profiler markers in
run_model()to mark prefill/decode phases, model forward passes, and CUDA graph replays - Added top-level profiler markers in
forward()to mark the entire forward pass phase
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
atom/model_engine/model_runner.py
Outdated
|
|
||
| with torch_profiler.record_function(f"forward_pass_{phase_marker}"): |
There was a problem hiding this comment.
The profiler marker naming is inconsistent between forward() and run_model(). In forward(), the markers are named "forward_pass_PREFILL" and "forward_pass_DECODE", while in run_model() they use "PREFILL_PHASE" and "DECODE_PHASE" with additional suffixes like "_model_forward" or "_cudagraph_replay_bs{graph_bs}". Consider using a consistent naming scheme across both methods for easier trace analysis, such as using "PHASE" suffix consistently or removing it from run_model().
| with torch_profiler.record_function(f"forward_pass_{phase_marker}"): | |
| with torch_profiler.record_function(f"{phase_marker}_PHASE"): |
|
thanks your contribution, @HaonanWang98 had a more powerful work on the way, i would l like that one |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 1 out of 1 changed files in this pull request and generated 2 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| logits = self.run_model(input_ids) | ||
| # Add top-level marker for the entire forward pass | ||
| is_prefill = batch.total_tokens_num_prefill > 0 | ||
| phase_name = "PREFILL" if is_prefill else "DECODE" |
There was a problem hiding this comment.
Trailing whitespace detected at the end of this line. Please remove the trailing spaces after the variable assignment.
| phase_name = "PREFILL" if is_prefill else "DECODE" | |
| phase_name = "PREFILL" if is_prefill else "DECODE" |
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist