Skip to content
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

Split the tensor allocate operation of comm tensors from function prepComms to function generate_io_tensors #173

Closed
wants to merge 2 commits into from
Closed

Split the tensor allocate operation of comm tensors from function prepComms to function generate_io_tensors #173

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3a84d81
split the tensor allocate operation of comm tensors from function pre…
Aug 1, 2024
610b035
Merge branch 'main' into shengf/pr-1
shengfukevin Aug 20, 2024
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
44 changes: 26 additions & 18 deletions et_replay/tools/comm_replay.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
import logging
import os
import time
from typing import Dict, List, Set
from typing import Dict, List, Set, Tuple

import numpy as np
import torch
Expand Down Expand Up @@ -610,12 +610,35 @@ def hashEtCommsOp(self, commsOp: commsArgs) -> int:

return hash(op)

def generate_io_tensors(
self,
curComm: commsArgs,
commsParams: commsParamsHolderBase,
regenerateTensors: bool
) -> Tuple[torch.Tensor, torch.Tensor]:
# Use exactly specified inMsgSize/outMsgSize if call from trace replay
# This avoid regenerating sizes such as in _prep_all_gather_base
commsParams.size_from_trace = True
commsParams.dtype = self.dtypeMap[curComm.dtype]
if not curComm.id or regenerateTensors:
return super().prepComm(curComm, commsParams)
else:
commsOpHash = self.hashEtCommsOp(curComm)
if commsOpHash in self.et_to_tensors:
# Allocate input/output tensors if first time replay, otherwise the previous ones.
super().prepComm(curComm, commsParams, False)
(ipTensor, opTensor) = self.et_to_tensors[commsOpHash]
else:
(ipTensor, opTensor) = super().prepComm(curComm, commsParams, True)
self.et_to_tensors[commsOpHash] = (ipTensor, opTensor)
return (ipTensor, opTensor)

def prepComms(
self,
curComm: commsArgs,
commsParams: commsParamsHolderBase,
regenerateTensors: bool = True,
) -> tuple[torch.Tensor, torch.Tensor]:
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Prepares the appropriate tensors for the current collective communication.

Expand Down Expand Up @@ -686,22 +709,7 @@ def prepComms(
f"shrink message sizes to curInNumElem {curComm.inMsgSize}, curOutNumElem {curComm.outMsgSize}"
)

# Use exactly specified inMsgSize/outMsgSize if call from trace replay
# This avoid regenerating sizes such as in _prep_all_gather_base
commsParams.size_from_trace = True
commsParams.dtype = self.dtypeMap[curComm.dtype]
if not curComm.id or regenerateTensors:
return super().prepComm(curComm, commsParams)
else:
commsOpHash = self.hashEtCommsOp(curComm)
if commsOpHash in self.et_to_tensors:
# Allocate input/output tensors if first time replay, otherwise the previous ones.
super().prepComm(curComm, commsParams, False)
(ipTensor, opTensor) = self.et_to_tensors[commsOpHash]
else:
(ipTensor, opTensor) = super().prepComm(curComm, commsParams, True)
self.et_to_tensors[commsOpHash] = (ipTensor, opTensor)
return (ipTensor, opTensor)
return self.generate_io_tensors(curComm, commsParams, regenerateTensors)

def commRebalance(self, curComm: commsArgs) -> None:
"""
Expand Down
Loading