Skip to content

Regenerate pt2_et trace with corrected stream_id, fix test #207

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

Closed
leloleo wants to merge 1 commit into from
Closed

Regenerate pt2_et trace with corrected stream_id, fix test #207

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: 9 additions & 4 deletions et_replay/tools/et_replay.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@
from param_bench.train.compute.python.lib import pytorch as lib_pytorch
from param_bench.train.compute.python.lib.init_helper import load_modules
from param_bench.train.compute.python.workloads import pytorch as workloads_pytorch
from torch._C import _cuda_getCurrentRawStream as get_raw_stream
from torch._inductor.async_compile import AsyncCompile

# grid and split_scan_grid are dynamically loaded
Expand Down Expand Up @@ -1457,10 +1458,14 @@ def run_op(self, node, iter, cnt): # noqa: C901
outputs = []
if output_count == 0:
if node.kernel_backend == "triton":
# The last entry in inputs is stream
# ET captured the raw pointer of cudaStream_t,
# but triton needs the stream object, hard code to 0 for now
exec("func.run(*inputs[:-1], stream=0)")
# The last entry in inputs is the stream id, however if we pass that back into func.run()
# ET will attemp to dereference the stream again. To fully support triton kernels we need
# to keep a mapping between stream address and stream id. This is not implemented yet so
# we will capture the raw stream from the device and pass that back to the replay.
# In reality all PT2 compute ops are in the same stream to begin with, only collectives
# are in different streams, so this is fine.
stream = get_raw_stream(self.cuda_id)
func.run(*inputs[:-1], stream=stream)
else:
func(*inputs)
else:
Expand Down