Repository: pytorch
Issue: [torch.compile][Megatron] Error with Megatron with Pytorch v2.5.0 using `AOTAutograd` and `torch.compile`
Labels: triaged,oncall: pt2,module: aotdispatch,module: pt2-dispatcher

### 🐛 Describe the bug

I am using  `torch.compile` with `AOTAutograd` to compile a Megatron model. It worked fine with Python v2.4.1, but after upgrading to Pytorch 2.5.0, I encountered the following error, but it was fine when running with `torch.compile` alone.

```
[rank0]: Traceback (most recent call last):
[rank0]:   File "/workspace/experiment/examples/megatron/megatron_test.py", line 117, in <module>
[rank0]:     main()
[rank0]:   File "/workspace/experiment/examples/megatron/megatron_test.py", line 69, in main
[rank0]:     compiler.compile(
[rank0]:   File "/workspace/experiment/experiment/compiler.py", line 289, in compile
[rank0]:     compiled_func = compiler.compile(module, loss_func, *args, **kwargs)
[rank0]:                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/workspace/experiment/experiment/compiler.py", line 225, in compile
[rank0]:     _ = compiled_func(*args, **kwargs)
[rank0]:         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 465, in _fn
[rank0]:     return fn(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/workspace/experiment/experiment/compiler.py", line 162, in step_func
[rank0]:     output_tensor = module(*args, **kwargs)
[rank0]:                     ^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/workspace/megatron-lm/megatron/legacy/model/module.py", line 189, in forward
[rank0]:     outputs = self.module(*inputs, **kwargs)
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1747, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/workspace/megatron-lm/megatron/core/models/gpt/gpt_model.py", line 238, in forward
[rank0]:     hidden_states = self.decoder(
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py", line 1333, in __call__
[rank0]:     return self._torchdynamo_orig_callable(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py", line 1124, in __call__
[rank0]:     result = self._inner_convert(
[rank0]:              ^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py", line 528, in __call__
[rank0]:     return _compile(
[rank0]:            ^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py", line 948, in _compile
[rank0]:     guarded_code = compile_inner(code, one_graph, hooks, transform)
[rank0]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py", line 679, in compile_inner
[rank0]:     return _compile_inner(code, one_graph, hooks, transform)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_utils_internal.py", line 87, in wrapper_function
[rank0]:     return function(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py", line 712, in _compile_inner
[rank0]:     out_code = transform_code_object(code, transform)
[rank0]:                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/bytecode_transformation.py", line 1337, in transform_code_object
[rank0]:     transformations(instructions, code_options)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py", line 221, in _fn
[rank0]:     return fn(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py", line 641, in transform
[rank0]:     tracer.run()
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/symbolic_convert.py", line 2766, in run
[rank0]:     super().run()
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/symbolic_convert.py", line 973, in run
[rank0]:     while self.step():
[rank0]:           ^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/symbolic_convert.py", line 885, in step
[rank0]:     self.dispatch_table[inst.opcode](self, inst)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/symbolic_convert.py", line 2957, in RETURN_VALUE
[rank0]:     self._return(inst)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/symbolic_convert.py", line 2942, in _return
[rank0]:     self.output.compile_subgraph(
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/output_graph.py", line 1142, in compile_subgraph
[rank0]:     self.compile_and_call_fx_graph(tx, pass2.graph_output_vars(), root)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/output_graph.py", line 1369, in compile_and_call_fx_graph
[rank0]:     compiled_fn = self.call_user_compiler(gm)
[rank0]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/output_graph.py", line 1416, in call_user_compiler
[rank0]:     return self._call_user_compiler(gm)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/output_graph.py", line 1465, in _call_user_compiler
[rank0]:     raise BackendCompilerFailed(self.compiler_fn, e).with_traceback(
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/output_graph.py", line 1446, in _call_user_compiler
[rank0]:     compiled_fn = compiler_fn(gm, self.example_inputs())
[rank0]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/repro/after_dynamo.py", line 130, in __call__
[rank0]:     compiled_gm = compiler_fn(gm, example_inputs)
[rank0]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/__init__.py", line 2280, in __call__
[rank0]:     return self.compiler_fn(model_, inputs_, **self.kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/workspace/experiment/experiment/compiler.py", line 212, in backend_func
[rank0]:     return aot_autograd(
[rank0]:            ^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/backends/common.py", line 72, in __call__
[rank0]:     cg = aot_module_simplified(gm, example_inputs, **self.kwargs)
[rank0]:          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_functorch/aot_autograd.py", line 1077, in aot_module_simplified
[rank0]:     compiled_fn = dispatch_and_compile()
[rank0]:                   ^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_functorch/aot_autograd.py", line 1062, in dispatch_and_compile
[rank0]:     compiled_fn, _ = create_aot_dispatcher_function(
[rank0]:                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_functorch/aot_autograd.py", line 523, in create_aot_dispatcher_function
[rank0]:     return _create_aot_dispatcher_function(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_functorch/aot_autograd.py", line 761, in _create_aot_dispatcher_function
[rank0]:     compiled_fn, fw_metadata = compiler_fn(
[rank0]:                                ^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 628, in aot_dispatch_autograd
[rank0]:     placeholder_list = fx_placeholder_vals(bw_module)
[rank0]:                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/symbolic_shapes.py", line 1213, in fx_placeholder_vals
[rank0]:     return [n.meta["val"] for n in gm.graph.nodes if n.op == "placeholder"]
[rank0]:             ~~~~~~^^^^^^^
[rank0]: torch._dynamo.exc.BackendCompilerFailed: backend='backend_func' raised:
[rank0]: KeyError: 'val'
```

It seems that the error is caused by the following function and the compilation works if I removed the line `predicted_logits[target_mask] = 0.0`.

```
    @staticmethod
    def calculate_predicted_logits(
        vocab_parallel_logits: torch.Tensor,
        target: torch.Tensor,
        logits_max: torch.Tensor,
        vocab_start_index: int,
        vocab_end_index: int,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Calculates predicted logits."""

        # In-place subtraction reduces memory pressure.
        vocab_parallel_logits -= logits_max.unsqueeze(dim=-1)

        # Create a mask of valid vocab ids (1 means it needs to be masked).
        target_mask = (target < vocab_start_index) | (target >= vocab_end_index)
        masked_target = target.clone() - vocab_start_index
        masked_target[target_mask] = 0

        # Get predicted-logits = logits[target].
        # For Simplicity, we convert logits to a 2-D tensor with size
        # [*, partition-vocab-size] and target to a 1-D tensor of size [*].
        partition_vocab_size = vocab_parallel_logits.size()[-1]
        logits_2d = vocab_parallel_logits.view(-1, partition_vocab_size)
        masked_target_1d = masked_target.view(-1)
        arange_1d = torch.arange(start=0, end=logits_2d.size()[0], device=logits_2d.device)
        predicted_logits_1d = logits_2d[arange_1d, masked_target_1d]
        predicted_logits_1d = predicted_logits_1d.clone().contiguous()
        predicted_logits = predicted_logits_1d.view_as(target)
        predicted_logits[target_mask] = 0.0

        exp_logits = vocab_parallel_logits
        # torch.exp(vocab_parallel_logits, out=exp_logits)
        exp_logits = torch.exp_(vocab_parallel_logits) # This is more dynamo-friendly
        sum_exp_logits = exp_logits.sum(dim=-1)

        return target_mask, masked_target_1d, predicted_logits, sum_exp_logits, exp_logits
```
The following is the generated graphs of the function, and I found that `_tensor_constant1` has no meta['val'].

```
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs] TRACED GRAPH
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]  ===== Forward graph 10 =====
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]  /mnt/data/mksit/anaconda3/envs/experiment/lib/python3.10/site-packages/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]     def forward(self, primals_1: "bf16[64, 8, 128][1024, 128, 1]cuda:0", primals_2: "bf16[100, 128][128, 1]cuda:0", primals_3: "i64[8, 64][64, 1]cuda:0"):
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]          # File: /home/mankit/workspace/experiment/third_party/Megatron-LM/megatron/core/tensor_parallel/mappings.py:526 in copy_to_tensor_model_parallel_region, code: return _CopyToModelParallelRegion.apply(input_)
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         view: "bf16[64, 8, 128][1024, 128, 1]cuda:0" = torch.ops.aten.view.default(primals_1, [64, 8, 128]);  primals_1 = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]          # File: /home/mankit/workspace/experiment/third_party/Megatron-LM/megatron/core/tensor_parallel/layers.py:670 in linear_with_grad_accumulation_and_async_allreduce, code: return LinearWithGradAccumulationAndAsyncCommunication.apply(*args)
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         t: "bf16[128, 100][1, 128]cuda:0" = torch.ops.aten.t.default(primals_2)
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         view_1: "bf16[512, 128][128, 1]cuda:0" = torch.ops.aten.view.default(view, [512, 128])
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         mm: "bf16[512, 100][100, 1]cuda:0" = torch.ops.aten.mm.default(view_1, t);  view_1 = t = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         view_2: "bf16[64, 8, 100][800, 100, 1]cuda:0" = torch.ops.aten.view.default(mm, [64, 8, 100]);  mm = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]          # File: /home/mankit/workspace/experiment/third_party/Megatron-LM/megatron/core/models/common/language_module/language_module.py:37 in compute_language_model_loss, code: labels = labels.transpose(0, 1).contiguous()
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         transpose: "i64[64, 8][1, 64]cuda:0" = torch.ops.aten.transpose.int(primals_3, 0, 1);  primals_3 = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         clone: "i64[64, 8][8, 1]cuda:0" = torch.ops.aten.clone.default(transpose, memory_format = torch.contiguous_format);  transpose = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]          # File: /home/mankit/workspace/experiment/third_party/Megatron-LM/megatron/core/tensor_parallel/cross_entropy.py:236 in vocab_parallel_cross_entropy, code: return _VocabParallelCrossEntropy.apply(vocab_parallel_logits, target, label_smoothing)
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         _to_copy: "f32[64, 8, 100][800, 100, 1]cuda:0" = torch.ops.aten._to_copy.default(view_2, dtype = torch.float32);  view_2 = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         max_1 = torch.ops.aten.max.dim(_to_copy, -1)
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         getitem: "f32[64, 8][8, 1]cuda:0" = max_1[0];  max_1 = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         all_reduce: "f32[64, 8][8, 1]cuda:0" = torch.ops._c10d_functional.all_reduce.default(getitem, 'max', '8')
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         wait_tensor: "f32[64, 8][8, 1]cuda:0" = torch.ops._c10d_functional.wait_tensor.default(all_reduce);  all_reduce = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         copy: "f32[64, 8][8, 1]cuda:0" = torch.ops.aten.copy.default(getitem, wait_tensor);  wait_tensor = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         unsqueeze_1: "f32[64, 8, 1][8, 1, 1]cuda:0" = torch.ops.aten.unsqueeze.default(copy, -1);  copy = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         sub: "f32[64, 8, 100][800, 100, 1]cuda:0" = torch.ops.aten.sub.Tensor(_to_copy, unsqueeze_1);  _to_copy = unsqueeze_1 = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         lt: "b8[64, 8][8, 1]cuda:0" = torch.ops.aten.lt.Scalar(clone, 0)
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         ge: "b8[64, 8][8, 1]cuda:0" = torch.ops.aten.ge.Scalar(clone, 100)
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         bitwise_or: "b8[64, 8][8, 1]cuda:0" = torch.ops.aten.bitwise_or.Tensor(lt, ge);  lt = ge = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         clone_1: "i64[64, 8][8, 1]cuda:0" = torch.ops.aten.clone.default(clone);  clone = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         sub_1: "i64[64, 8][8, 1]cuda:0" = torch.ops.aten.sub.Tensor(clone_1, 0);  clone_1 = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         _tensor_constant0 = self._tensor_constant0
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         lift_fresh_copy: "i64[][]cuda:0" = torch.ops.aten.lift_fresh_copy.default(_tensor_constant0);  _tensor_constant0 = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         where: "i64[64, 8][8, 1]cuda:0" = torch.ops.aten.where.self(bitwise_or, lift_fresh_copy, sub_1);  lift_fresh_copy = sub_1 = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         view_4: "i64[512][1]cuda:0" = torch.ops.aten.view.default(where, [-1]);  where = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         arange: "i64[512][1]cuda:0" = torch.ops.aten.arange.start(0, 512, device = device(type='cuda', index=0), pin_memory = False)
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         view_5: "f32[512, 100][100, 1]cuda:0" = torch.ops.aten.view.default(sub, [-1, 100])
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         index: "f32[512][1]cuda:0" = torch.ops.aten.index.Tensor(view_5, [arange, view_4]);  view_5 = arange = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         clone_2: "f32[512][1]cuda:0" = torch.ops.aten.clone.default(index);  index = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         view_6: "f32[64, 8][8, 1]cuda:0" = torch.ops.aten.view.default(clone_2, [64, 8]);  clone_2 = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         _tensor_constant1 = self._tensor_constant1
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         lift_fresh_copy_1: "f32[][]cuda:0" = torch.ops.aten.lift_fresh_copy.default(_tensor_constant1)
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         where_1: "f32[64, 8][8, 1]cuda:0" = torch.ops.aten.where.self(bitwise_or, lift_fresh_copy_1, view_6);  lift_fresh_copy_1 = view_6 = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         exp: "f32[64, 8, 100][800, 100, 1]cuda:0" = torch.ops.aten.exp.default(sub);  sub = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         sum_1: "f32[64, 8][8, 1]cuda:0" = torch.ops.aten.sum.dim_IntList(exp, [-1])
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         all_reduce_1: "f32[64, 8][8, 1]cuda:0" = torch.ops._c10d_functional.all_reduce.default(where_1, 'sum', '8')
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         wait_tensor_1: "f32[64, 8][8, 1]cuda:0" = torch.ops._c10d_functional.wait_tensor.default(all_reduce_1);  all_reduce_1 = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         copy_1: "f32[64, 8][8, 1]cuda:0" = torch.ops.aten.copy.default(where_1, wait_tensor_1);  where_1 = wait_tensor_1 = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         all_reduce_2: "f32[64, 8][8, 1]cuda:0" = torch.ops._c10d_functional.all_reduce.default(sum_1, 'sum', '8')
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         wait_tensor_2: "f32[64, 8][8, 1]cuda:0" = torch.ops._c10d_functional.wait_tensor.default(all_reduce_2);  all_reduce_2 = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         copy_2: "f32[64, 8][8, 1]cuda:0" = torch.ops.aten.copy.default(sum_1, wait_tensor_2);  sum_1 = wait_tensor_2 = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         log: "f32[64, 8][8, 1]cuda:0" = torch.ops.aten.log.default(copy_2)
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         sub_2: "f32[64, 8][8, 1]cuda:0" = torch.ops.aten.sub.Tensor(log, copy_1);  log = copy_1 = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]          # File: /home/mankit/workspace/experiment/third_party/Megatron-LM/megatron/core/models/common/language_module/language_module.py:44 in compute_language_model_loss, code: loss = loss.transpose(0, 1).contiguous()
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         transpose_1: "f32[8, 64][1, 8]cuda:0" = torch.ops.aten.transpose.int(sub_2, 0, 1);  sub_2 = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         clone_3: "f32[8, 64][64, 1]cuda:0" = torch.ops.aten.clone.default(transpose_1, memory_format = torch.contiguous_format);  transpose_1 = None
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         return (clone_3, primals_2, view, getitem, bitwise_or, view_4, _tensor_constant1, exp, copy_2)
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs]         
[rank0]:I1128 20:51:42.850000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:523] [25/0] [__aot_graphs] 
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs] TRACED GRAPH
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]  ===== Backward graph 10 =====
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]  /mnt/data/mksit/anaconda3/envs/experiment/lib/python3.10/site-packages/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]     def forward(self, primals_2: "bf16[100, 128][128, 1]cuda:0", view: "bf16[64, 8, 128][1024, 128, 1]cuda:0", getitem: "f32[64, 8][8, 1]cuda:0", bitwise_or: "b8[64, 8][8, 1]cuda:0", view_4: "i64[512][1]cuda:0", _tensor_constant1, exp: "f32[64, 8, 100][800, 100, 1]cuda:0", copy_2: "f32[64, 8][8, 1]cuda:0", tangents_1: "f32[8, 64][64, 1]cuda:0"):
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]          # File: /home/mankit/workspace/experiment/third_party/Megatron-LM/megatron/core/tensor_parallel/layers.py:670 in linear_with_grad_accumulation_and_async_allreduce, code: return LinearWithGradAccumulationAndAsyncCommunication.apply(*args)
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         t: "bf16[128, 100][1, 128]cuda:0" = torch.ops.aten.t.default(primals_2)
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         view_1: "bf16[512, 128][128, 1]cuda:0" = torch.ops.aten.view.default(view, [512, 128])
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         mm: "bf16[512, 100][100, 1]cuda:0" = torch.ops.aten.mm.default(view_1, t);  view_1 = t = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         view_2: "bf16[64, 8, 100][800, 100, 1]cuda:0" = torch.ops.aten.view.default(mm, [64, 8, 100]);  mm = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]          # File: /home/mankit/workspace/experiment/third_party/Megatron-LM/megatron/core/tensor_parallel/cross_entropy.py:236 in vocab_parallel_cross_entropy, code: return _VocabParallelCrossEntropy.apply(vocab_parallel_logits, target, label_smoothing)
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         _to_copy: "f32[64, 8, 100][800, 100, 1]cuda:0" = torch.ops.aten._to_copy.default(view_2, dtype = torch.float32);  view_2 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         all_reduce: "f32[64, 8][8, 1]cuda:0" = torch.ops._c10d_functional.all_reduce.default(getitem, 'max', '8')
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         wait_tensor: "f32[64, 8][8, 1]cuda:0" = torch.ops._c10d_functional.wait_tensor.default(all_reduce);  all_reduce = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         copy: "f32[64, 8][8, 1]cuda:0" = torch.ops.aten.copy.default(getitem, wait_tensor);  getitem = wait_tensor = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         unsqueeze_1: "f32[64, 8, 1][8, 1, 1]cuda:0" = torch.ops.aten.unsqueeze.default(copy, -1);  copy = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         sub: "f32[64, 8, 100][800, 100, 1]cuda:0" = torch.ops.aten.sub.Tensor(_to_copy, unsqueeze_1);  _to_copy = unsqueeze_1 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         arange: "i64[512][1]cuda:0" = torch.ops.aten.arange.start(0, 512, device = device(type='cuda', index=0), pin_memory = False)
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         view_5: "f32[512, 100][100, 1]cuda:0" = torch.ops.aten.view.default(sub, [-1, 100]);  sub = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         index: "f32[512][1]cuda:0" = torch.ops.aten.index.Tensor(view_5, [arange, view_4]);  view_5 = arange = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         clone_2: "f32[512][1]cuda:0" = torch.ops.aten.clone.default(index);  index = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         view_6: "f32[64, 8][8, 1]cuda:0" = torch.ops.aten.view.default(clone_2, [64, 8]);  clone_2 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         lift_fresh_copy_1: "f32[][]cuda:0" = torch.ops.aten.lift_fresh_copy.default(_tensor_constant1);  _tensor_constant1 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         where_1: "f32[64, 8][8, 1]cuda:0" = torch.ops.aten.where.self(bitwise_or, lift_fresh_copy_1, view_6);  lift_fresh_copy_1 = view_6 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         sum_1: "f32[64, 8][8, 1]cuda:0" = torch.ops.aten.sum.dim_IntList(exp, [-1])
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         all_reduce_1: "f32[64, 8][8, 1]cuda:0" = torch.ops._c10d_functional.all_reduce.default(where_1, 'sum', '8');  where_1 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         wait_tensor_1: "f32[64, 8][8, 1]cuda:0" = torch.ops._c10d_functional.wait_tensor.default(all_reduce_1);  all_reduce_1 = wait_tensor_1 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         all_reduce_2: "f32[64, 8][8, 1]cuda:0" = torch.ops._c10d_functional.all_reduce.default(sum_1, 'sum', '8');  sum_1 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         wait_tensor_2: "f32[64, 8][8, 1]cuda:0" = torch.ops._c10d_functional.wait_tensor.default(all_reduce_2);  all_reduce_2 = wait_tensor_2 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         unsqueeze_3: "f32[64, 8, 1][8, 1, 1]cuda:0" = torch.ops.aten.unsqueeze.default(copy_2, -1);  copy_2 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         div: "f32[64, 8, 100][800, 100, 1]cuda:0" = torch.ops.aten.div.Tensor(exp, unsqueeze_3);  exp = unsqueeze_3 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]          # File: /home/mankit/workspace/experiment/third_party/Megatron-LM/megatron/core/models/common/language_module/language_module.py:44 in compute_language_model_loss, code: loss = loss.transpose(0, 1).contiguous()
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         transpose_2: "f32[64, 8][1, 64]cuda:0" = torch.ops.aten.transpose.int(tangents_1, 0, 1);  tangents_1 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]          # File: /home/mankit/workspace/experiment/third_party/Megatron-LM/megatron/core/tensor_parallel/cross_entropy.py:236 in vocab_parallel_cross_entropy, code: return _VocabParallelCrossEntropy.apply(vocab_parallel_logits, target, label_smoothing)
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         view_7: "f32[512, 100][100, 1]cuda:0" = torch.ops.aten.view.default(div, [-1, 100]);  div = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         arange_1: "i64[512][1]cuda:0" = torch.ops.aten.arange.start(0, 512, device = device(type='cuda', index=0), pin_memory = False)
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         view_8: "b8[512][1]cuda:0" = torch.ops.aten.view.default(bitwise_or, [-1]);  bitwise_or = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         _to_copy_1: "f32[512][1]cuda:0" = torch.ops.aten._to_copy.default(view_8, dtype = torch.float32);  view_8 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         rsub: "f32[512][1]cuda:0" = torch.ops.aten.rsub.Scalar(_to_copy_1, 1.0);  _to_copy_1 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         index_1: "f32[512][1]cuda:0" = torch.ops.aten.index.Tensor(view_7, [arange_1, view_4])
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         sub_3: "f32[512][1]cuda:0" = torch.ops.aten.sub.Tensor(index_1, rsub);  index_1 = rsub = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         index_put: "f32[512, 100][100, 1]cuda:0" = torch.ops.aten.index_put.default(view_7, [arange_1, view_4], sub_3);  view_7 = arange_1 = view_4 = sub_3 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         view_9: "f32[64, 8, 100][800, 100, 1]cuda:0" = torch.ops.aten.view.default(index_put, [64, 8, 100]);  index_put = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         unsqueeze_4: "f32[64, 8, 1][1, 64, 1]cuda:0" = torch.ops.aten.unsqueeze.default(transpose_2, -1);  transpose_2 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         mul: "f32[64, 8, 100][800, 100, 1]cuda:0" = torch.ops.aten.mul.Tensor(view_9, unsqueeze_4);  view_9 = unsqueeze_4 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         _to_copy_2: "bf16[64, 8, 100][800, 100, 1]cuda:0" = torch.ops.aten._to_copy.default(mul, dtype = torch.bfloat16);  mul = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]          # File: /home/mankit/workspace/experiment/third_party/Megatron-LM/megatron/core/tensor_parallel/layers.py:670 in linear_with_grad_accumulation_and_async_allreduce, code: return LinearWithGradAccumulationAndAsyncCommunication.apply(*args)
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         view_11: "bf16[512, 100][100, 1]cuda:0" = torch.ops.aten.view.default(_to_copy_2, [512, 100])
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         mm_1: "bf16[512, 128][128, 1]cuda:0" = torch.ops.aten.mm.default(view_11, primals_2);  view_11 = primals_2 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         view_12: "bf16[64, 8, 128][1024, 128, 1]cuda:0" = torch.ops.aten.view.default(mm_1, [64, 8, 128]);  mm_1 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         view_13: "bf16[512, 100][100, 1]cuda:0" = torch.ops.aten.view.default(_to_copy_2, [512, 100]);  _to_copy_2 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         view_14: "bf16[512, 128][128, 1]cuda:0" = torch.ops.aten.view.default(view, [512, 128]);  view = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         t_1: "bf16[100, 512][1, 100]cuda:0" = torch.ops.aten.t.default(view_13);  view_13 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         mm_2: "bf16[100, 128][128, 1]cuda:0" = torch.ops.aten.mm.default(t_1, view_14);  t_1 = view_14 = None
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         return (view_12, mm_2, None)
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs]         
[rank0]:I1128 20:51:42.852000 1695611 site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:534] [25/0] [__aot_graphs] 
```




### Versions

Collecting environment information...
PyTorch version: 2.5.0+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 12.3.0-1ubuntu1~22.04) 12.3.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.10.15 (main, Oct  3 2024, 07:27:34) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-6.8.0-40-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.6.68
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA RTX A5000
GPU 1: NVIDIA RTX A5000
GPU 2: NVIDIA RTX A5000
GPU 3: NVIDIA RTX A5000
GPU 4: NVIDIA RTX A5000
GPU 5: NVIDIA RTX A5000
GPU 6: NVIDIA RTX A5000
GPU 7: NVIDIA RTX A5000

Nvidia driver version: 560.35.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.7
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.7
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.7
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.7
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.7
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.7
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.7
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                         x86_64
CPU op-mode(s):                       32-bit, 64-bit
Address sizes:                        43 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               56
On-line CPU(s) list:                  0-55
Vendor ID:                            AuthenticAMD
Model name:                           AMD EPYC 7453 28-Core Processor
CPU family:                           25
Model:                                1
Thread(s) per core:                   1
Core(s) per socket:                   28
Socket(s):                            2
Stepping:                             1
Frequency boost:                      enabled
CPU max MHz:                          3488.5249
CPU min MHz:                          1500.0000
BogoMIPS:                             5499.46
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local user_shstk clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin brs arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif v_spec_ctrl umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca sev sev_es debug_swap
Virtualisation:                       AMD-V
L1d cache:                            1.8 MiB (56 instances)
L1i cache:                            1.8 MiB (56 instances)
L2 cache:                             28 MiB (56 instances)
L3 cache:                             128 MiB (8 instances)
NUMA node(s):                         2
NUMA node0 CPU(s):                    0-27
NUMA node1 CPU(s):                    28-55
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Not affected
Vulnerability Spec rstack overflow:   Mitigation; Safe RET
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:             Mitigation; Retpolines; IBPB conditional; IBRS_FW; STIBP disabled; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] mypy==1.13.0
[pip3] mypy-extensions==1.0.0
[pip3] mypy-protobuf==3.6.0
[pip3] numpy==1.26.3
[pip3] nvidia-cublas-cu12==12.4.5.8
[pip3] nvidia-cuda-cupti-cu12==12.4.127
[pip3] nvidia-cuda-nvrtc-cu12==12.4.127
[pip3] nvidia-cuda-runtime-cu12==12.4.127
[pip3] nvidia-cudnn-cu12==9.1.0.70
[pip3] nvidia-cufft-cu12==11.2.1.3
[pip3] nvidia-curand-cu12==10.3.5.147
[pip3] nvidia-cusolver-cu12==11.6.1.9
[pip3] nvidia-cusparse-cu12==12.3.1.170
[pip3] nvidia-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu12==12.4.127
[pip3] nvidia-nvtx-cu12==12.4.127
[pip3] torch==2.5.0+cu124
[pip3] torchaudio==2.5.0+cu124
[pip3] torchvision==0.20.0+cu124
[pip3] torchviz==0.0.2
[pip3] triton==3.1.0
[conda] numpy                     1.26.3                   pypi_0    pypi
[conda] nvidia-cublas-cu12        12.4.5.8                 pypi_0    pypi
[conda] nvidia-cuda-cupti-cu12    12.4.127                 pypi_0    pypi
[conda] nvidia-cuda-nvrtc-cu12    12.4.127                 pypi_0    pypi
[conda] nvidia-cuda-runtime-cu12  12.4.127                 pypi_0    pypi
[conda] nvidia-cudnn-cu12         9.1.0.70                 pypi_0    pypi
[conda] nvidia-cufft-cu12         11.2.1.3                 pypi_0    pypi
[conda] nvidia-curand-cu12        10.3.5.147               pypi_0    pypi
[conda] nvidia-cusolver-cu12      11.6.1.9                 pypi_0    pypi
[conda] nvidia-cusparse-cu12      12.3.1.170               pypi_0    pypi
[conda] nvidia-nccl-cu12          2.21.5                   pypi_0    pypi
[conda] nvidia-nvjitlink-cu12     12.4.127                 pypi_0    pypi
[conda] nvidia-nvtx-cu12          12.4.127                 pypi_0    pypi
[conda] torch                     2.5.0+cu124              pypi_0    pypi
[conda] torchaudio                2.5.0+cu124              pypi_0    pypi
[conda] torchvision               0.20.0+cu124             pypi_0    pypi
[conda] torchviz                  0.0.2                    pypi_0    pypi
[conda] triton                    3.1.0                    pypi_0    pypi

cc @chauhang @penguinwu @zou3519 @bdhirsh @yf225