GitLab

[fx] modify the calculation of node_size in MetaInfoProp for activation checkpointing usages (#1425)

* [fx] modify the calculation of node_size in MetaInfoProp for activation checkpointing usages

* [fx] modify the calculation of node_size in MetaInfoProp for activation checkpointing usages

* [fx] modify the calculation of node_size in MetaInfoProp for activation checkpointing usages

* [tensor] add shape consistency feature to supportauto sharding spec transform.

* [tensor] remove unused argument in simulator, add doc string for target pair.

* support p2p communication with any type of object | pass test

* reconstruct pipeline schedule with p2p_v2.py(support communication with List[Any]) | pass test

* [communication] add p2p_v2.py to support communication with List[Any]

* Delete _pipeline_schedule_v2.py

* Delete test_cifar_with_data_pipeline_tensor_v2.py

* [engin/schedule] use p2p_v2 to recontruct pipeline_schedule

* [communication] remove print code

* [communication] remove print code

Co-authored-by: github-actions <github-actions@github.com>

Co-authored-by: github-actions <github-actions@github.com>

Co-authored-by: github-actions <github-actions@github.com>

* [device] add DeviceMesh class to support logical device layout

* polish code

* add doc string

* [fx] patched torch.max and data movement operator

* polish code

* zero optim state_dict takes only_rank_0

* fix unit test

Co-authored-by: github-actions <github-actions@github.com>

* fix zero ddp buffer cast

* fix zero ddp ignore params