saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-25 19:48:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[ORTModule] Use Env Variable to Set Provider Option cudnn_conv_algo_search (#13296)

Browse source 
This PR is to add support of using env variable to set provider option
cudnn_conv_algo_search so that user can choose better conv algo search
method to run model. This is a quick fix to unblock the test of MoE
model. Will have another PR to design and implement the ORTModule config
so that we can config ORTModule using Python script or config file
instead of env variable.
This commit is contained in:
Vincent Wang 2022-10-13 15:36:21 +08:00 committed by GitHub
parent 6fb70a82df
commit 807b2f4dd5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 28 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py
View file

@ -266,8 +266,12 @@ class GraphExecutionManager(GraphExecutionInterface):
providers.append("CPUExecutionProvider")
provider_option_map = {"device_id": str(self._device.index)}
if not self.is_rocm_pytorch:
# Set Conv algo search mode to HEURISTIC, which is same as PyTorch's default setting.
provider_option_map["cudnn_conv_algo_search"] = "HEURISTIC"
# Set Conv algo search mode to HEURISTIC by default, which is same as PyTorch's default setting.
conv_algo_search = ortmodule._defined_from_envvar("CONV_ALGO_SEARCH", "HEURISTIC", warn=True)
if conv_algo_search not in ["HEURISTIC", "EXHAUSTIVE"]:
warnings.warn("Invalid value of env CONV_ALGO_SEARCH. Must be HEURISTIC or EXHAUSTIVE.")
conv_algo_search = "HEURISTIC"
provider_option_map["cudnn_conv_algo_search"] = conv_algo_search
provider_option_map["cudnn_conv_use_max_workspace"] = "1"
provider_option_map["cudnn_conv1d_pad_to_nc1d"] = "1"
if self._use_external_gpu_allocator:

24
orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py
View file

@ -759,7 +759,8 @@ def test_scatternd_correctness(device, indices):
@pytest.mark.parametrize("use_fp16", [False, True])
@pytest.mark.parametrize("input_requires_grad", [False, True])
def test_gradient_correctness_conv1d(use_fp16, input_requires_grad):
@pytest.mark.parametrize("conv_algo_search", [None, "EXHAUSTIVE", "HEURISTIC"])
def test_gradient_correctness_conv1d(use_fp16, input_requires_grad, conv_algo_search):
class NeuralNetConv1D(torch.nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, padding=0, groups=1):
super(NeuralNetConv1D, self).__init__()
@ -775,6 +776,9 @@ def test_gradient_correctness_conv1d(use_fp16, input_requires_grad):
if torch.cuda.get_device_capability()[0] < 7:
return
if conv_algo_search is not None:
os.environ["CONV_ALGO_SEARCH"] = conv_algo_search
device = "cuda"
N, seq_len, C_in, C_out, kernel_size = 32, 128, 1536, 1536, 3
pt_model = NeuralNetConv1D(C_in, C_out, kernel_size, padding=1).to(device)
@ -787,12 +791,13 @@ def test_gradient_correctness_conv1d(use_fp16, input_requires_grad):
loss.backward()
return prediction
torch.manual_seed(2333)
for _ in range(10):
x = torch.randn(N, seq_len, C_in, device=device, requires_grad=input_requires_grad)
pt_prediction = run_step(pt_model, x)
ort_prediction = run_step(ort_model, x)
# PyTorch's Conv/GonvGrad uses HEURISTIC mode to search algo while ORT uses EXHAUSTIVE mode by default.
# PyTorch's Conv/GonvGrad uses HEURISTIC mode to search algo while this UT tests different modes for ORTModule.
# While different algo types generate slightly different results, especially for FP16,
# so relax the tolerance for comparison, especially for FP16 run and gradient comparison.
if use_fp16:
@ -802,6 +807,21 @@ def test_gradient_correctness_conv1d(use_fp16, input_requires_grad):
_test_helpers.assert_values_are_close(ort_prediction, pt_prediction, atol=1e-5)
_test_helpers.assert_gradients_match_and_reset_gradient(ort_model, pt_model, rtol=5e-2, atol=4e-2)
provider_options = ort_model._torch_module._execution_manager(
ort_model._is_training()
)._execution_agent._inference_session._provider_options
expected_conv_algo_search = "HEURISTIC" if conv_algo_search is None else conv_algo_search
actual_conv_algo_search = None
if "CUDAExecutionProvider" in provider_options:
actual_conv_algo_search = provider_options["CUDAExecutionProvider"]["cudnn_conv_algo_search"]
elif "ROCMExecutionProvider" in provider_options:
actual_conv_algo_search = provider_options["ROCMExecutionProvider"]["cudnn_conv_algo_search"]
assert actual_conv_algo_search == expected_conv_algo_search
if conv_algo_search is not None:
del os.environ["CONV_ALGO_SEARCH"]
def _run_gradient_correctness_transpose(perm, shape):
class NeuralNetTranspose(torch.nn.Module):