[Test] Unit test robustness

[comaniac]

Since the CI is set up., we now use this issue to track the progress of adding unit tests to make the repo more robustness. Specifically, we will cover the following checklist:

• Temporary end-to-end tests. Add a few end-to-end tests from examples/ to make sure the changes won't break existing benchmarks. ([Test] Add end-to-end tests #14)
• Unit tests - tracer. ([Test] Add tracer unit tests #16)
• Unit tests - initialization and consolidation. ([Test] Add default initialization test #54)
• Unit tests - schedule primitives
  • schedule hierarchy and sub-schedule accessing. ([Tracer] Add flatten argument to .trace() #29)
  • tracing. ([Tracer] Add flatten argument to .trace() #29)
  • sharding and syncing.
  • find node and subgraph pattern matching. ([Schedule] Refactor subgraph matching #35)
  • replace function and module. ([Schedule] Create subschedule for subgraph replacement #52)
  • operator fusion. ([Schedule] Add .fuse() primitive #25)
  • activation checkpointing (module and subgraph). ([Schedule] Support partial checkpointing #55)
  • pipeline partitioning.
• Unit tests - auto-tuning.
• Unit tests - DeepSpeed dialect.
• Remove temporary end-to-end tests.

Anyone could claim the above tasks by commenting this issue with an ETA. Please also link the submitted PRs to this issue for tracking.

[chhzh123]

Pipeline partition test is added in #103.

[chhzh123]

DeepSpeed dialect testing is added in #105.