OPT Finetune CreateStateParallel KeyError with PipeshardParallel

[zw123han]

Hello,

I'm tuning OPT using https://github.com/alpa-projects/alpa/tree/main/examples/opt_finetune with FollowParallel and CreateStateParallel to offload the peak CPU memory to my devices. This should also resolve issue #811

However, when method contains PipeshardParallel the following error is displayed. The same error is not observed when PipeshardParallel degenerates into ShardParallel (i.e. pipeline_parallel=1), nor with data_parallel + operator_parallel only.

To Reproduce

I attach my modified run_clm_flax.py (relevant lines: 725, 829). You can directly replace the version from https://github.com/alpa-projects/alpa/tree/main/examples/opt_finetune and reproduce the error with the existing run_2.7b_pipe.sh launch script.

https://github.com/zw123han/alpa/blob/main/examples/opt_finetune/run_clm_flax.py

---

Aside, do you think it would be helpful to merge this as the default OPT finetune script once stabilized to avoid CPU bottleneck for large models?

[merrymercy]

solved by #873.
See also #858