Support multi-gpu child ops and remove CPU/GPU op splitting

[TaekyungHeo]

Summary

• Introduced support for associating multiple GPU child operators with a single CPU operator.
• Added support for Tensor(c10::BFloat16) in get_data_type_size.
• Removed the method for splitting overlapping CPU and GPU operations due to its impact on operator sequencing.
  • Observed error rate reduction from 40% to 6% in Megatron traces by eliminating splitting logic.

Test Plan

$ python tools/trace_link.py --pytorch-et-file ~/Downloads/megatron_et_0.json --kineto-file ~/Downloads/megatron_kineto_0.json --output-file ~/rank0.json &
$ python tools/trace_link.py --pytorch-et-file ~/Downloads/megatron_et_1.json --kineto-file ~/Downloads/megatron_kineto_1.json --output-file ~/rank1.json &
$ python tools/trace_link.py --pytorch-et-file ~/Downloads/megatron_et_2.json --kineto-file ~/Downloads/megatron_kineto_2.json --output-file ~/rank2.json &
$ python tools/trace_link.py --pytorch-et-file ~/Downloads/megatron_et_3.json --kineto-file ~/Downloads/megatron_kineto_3.json --output-file ~/rank3.json &
$ python tools/trace_link.py --pytorch-et-file ~/Downloads/megatron_et_4.json --kineto-file ~/Downloads/megatron_kineto_4.json --output-file ~/rank4.json &
$ python tools/trace_link.py --pytorch-et-file ~/Downloads/megatron_et_5.json --kineto-file ~/Downloads/megatron_kineto_5.json --output-file ~/rank5.json &
$ python tools/trace_link.py --pytorch-et-file ~/Downloads/megatron_et_6.json --kineto-file ~/Downloads/megatron_kineto_6.json --output-file ~/rank6.json &
$ python tools/trace_link.py --pytorch-et-file ~/Downloads/megatron_et_7.json --kineto-file ~/Downloads/megatron_kineto_7.json --output-file ~/rank7.json &

$ python3 -m chakra.et_converter.et_converter --input_type PyTorch --input_filename ~/rank0.json --output_filename ~/rank.0.et --num_dims 1 &
$ python3 -m chakra.et_converter.et_converter --input_type PyTorch --input_filename ~/rank1.json --output_filename ~/rank.1.et --num_dims 1 &
$ python3 -m chakra.et_converter.et_converter --input_type PyTorch --input_filename ~/rank2.json --output_filename ~/rank.2.et --num_dims 1 &
$ python3 -m chakra.et_converter.et_converter --input_type PyTorch --input_filename ~/rank3.json --output_filename ~/rank.3.et --num_dims 1 &
$ python3 -m chakra.et_converter.et_converter --input_type PyTorch --input_filename ~/rank4.json --output_filename ~/rank.4.et --num_dims 1 &
$ python3 -m chakra.et_converter.et_converter --input_type PyTorch --input_filename ~/rank5.json --output_filename ~/rank.5.et --num_dims 1 &
$ python3 -m chakra.et_converter.et_converter --input_type PyTorch --input_filename ~/rank6.json --output_filename ~/rank.6.et --num_dims 1 &
$ python3 -m chakra.et_converter.et_converter --input_type PyTorch --input_filename ~/rank7.json --output_filename ~/rank.7.et --num_dims 1 &

• Profiled duration: 1345.336ms
• Simulated duration: 1431.253ms
• Error: 6.4%

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[srinivas212]

LGTM. Thanks for these improvements!