[QUESTION]Does Megatron support tracing computation graphs with torch.fx?

[fy-j]

I am trying to trace a computation graph in Megatron using torch.fx. However, I encountered the following error:

[rank1]: Traceback (most recent call last):
[rank1]:   File "/kimchou/GeeS***/python/gees***/adapters/pytorch/getTorchGraph.py", line 58, in getTorchGraph
[rank1]:     traced = symbolic_trace(model) # use transformers.utils.fx to trace the model
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/torch/fx/_symbolic_trace.py", line 1193, in symbolic_trace
[rank1]:     graph = tracer.trace(root, concrete_args)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/eval_frame.py", line 437, in _fn
[rank1]:     return fn(*args, **kwargs)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/external_utils.py", line 36, in inner
[rank1]:     return fn(*args, **kwargs)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/torch/fx/_symbolic_trace.py", line 793, in trace
[rank1]:     (self.create_arg(fn(*args)),),
[rank1]:   File "/kimchou/Megatron-LM-test/megatron/core/models/gpt/gpt_model.py", line 343, in forward
[rank1]:     **(extra_block_kwargs or {}),
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/torch/fx/proxy.py", line 443, in __bool__
[rank1]:     return self.tracer.to_bool(self)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/torch/fx/proxy.py", line 303, in to_bool
[rank1]:     raise TraceError('symbolically traced variables cannot be used as inputs to control flow')

How can I use torch.fx to trace the model graph