[Bugfix] Reproduce experimental results in docker image

[chhzh123]

Description

This PR tries to reproduce the experimental results presented in the paper using the provided docker image. Most of the errors are dependency issues.

Checklist

Following items list the failed test cases and the root causes:

• slapo-megatron, albert, gpu=1,2,4,8 (networkx is required by functorch, epoi must use the lastest commit to support the injection policy)
• slapo-deepspeed, albert, gpu=2,4,8 (same issue)
• slapo-megatron, gpt, gpu=2,4,8 (same issue)
• deepspeed, gpt, gpu=2,4,8 (missing datasets library)
• slapo-deepspeed, gpt, gpu=2,4,8 (same issue)
• slapo-megatron, wideresnet, gpu=2,4,8 (TorchScript is used in multi-device)
• slapo-deepspeed, wideresnet, gpu=2,4,8 (same issue)

[chhzh123]

Has run all the end-to-end single-device and single-node experiments in the docker and verified the results can be reproduced.

[comaniac]

I'll update the docker image accordingly.

[comaniac]

Docker image updated.

[comaniac]

Otherwise LGTM