Skip to content

GitLab

Unverified Commit ded8b9bd authored by Kirthi Shankar Sivamani's avatar Kirthi Shankar Sivamani Committed by GitHub
Browse files

Fix numerics for activation recompute (#327) 


Signed-off-by: default avatarKirthi Shankar Sivamani <ksivamani@nvidia.com>
parent 2a81e939
Hide whitespace changes
Inline Side-by-side
Showing with 4 additions and 3 deletions
transformer_engine/pytorch/jit.py
View file @ ded8b9bd
...@@ -220,7 +220,7 @@ def warmup_jit_bias_gelu( ...@@ -220,7 +220,7 @@ def warmup_jit_bias_gelu(
bias = torch.rand(ffn_hidden_size_per_partition, dtype=dtype, device="cuda") bias = torch.rand(ffn_hidden_size_per_partition, dtype=dtype, device="cuda")
inp = torch.rand( inp = torch.rand(
(seq_length, micro_batch_size, ffn_hidden_size_per_partition), (seq_length * micro_batch_size, ffn_hidden_size_per_partition),
dtype=dtype, dtype=dtype,
device="cuda", device="cuda",
) )
...@@ -229,8 +229,9 @@ def warmup_jit_bias_gelu( ...@@ -229,8 +229,9 @@ def warmup_jit_bias_gelu(
for bias_grad, input_grad in zip([True, True], [False, True]): for bias_grad, input_grad in zip([True, True], [False, True]):
bias.requires_grad, inp.requires_grad = bias_grad, input_grad bias.requires_grad, inp.requires_grad = bias_grad, input_grad
for _ in range(5): for _ in range(5):
output = bias_gelu_fused(inp, bias) _ = bias_gelu_fused_(inp, bias)
del bias, inp, output _ = gelu_fused_(inp)
del bias, inp
torch.cuda.empty_cache() torch.cuda.empty_cache()
torch.cuda.set_rng_state(rng_state) torch.cuda.set_rng_state(rng_state)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment