[Test] Unit test robustness

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:

- [x] Temporary end-to-end tests. Add a few end-to-end tests from `examples/` to make sure the changes won't break existing benchmarks. (#14)
- [x] Unit tests - tracer. (#16)
- [x] Unit tests - initialization and consolidation. (#54)
- [ ] Unit tests - schedule primitives
- [x] schedule hierarchy and sub-schedule accessing. (#29)
- [x] tracing. (#29)
- [x] sharding and syncing.
- [x] find node and subgraph pattern matching. (#35)
- [x] replace function and module. (#52)
- [x] operator fusion. (#25)
- [x] activation checkpointing (module and subgraph). (#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.