Add NVIDIA/CUDA checkpoint: config + remote code (from InternLumina-U2 GitHub, NPU-free)
ac8260b verified | from torch.distributed._tensor import Shard | |
| from veomni.distributed.parallel_plan import ParallelPlan | |
| def get_parallel_plan(): | |
| ep_plan = { | |
| "model.layers.*.mlp.experts.gate_proj": Shard(0), | |
| "model.layers.*.mlp.experts.up_proj": Shard(0), | |
| "model.layers.*.mlp.experts.down_proj": Shard(0), | |
| # Split export: experts is a ModuleList and each projection is a weight. | |
| "model.layers.*.mlp.experts.*.gate_proj.weight": Shard(0), | |
| "model.layers.*.mlp.experts.*.up_proj.weight": Shard(0), | |
| "model.layers.*.mlp.experts.*.down_proj.weight": Shard(0), | |
| } | |
| parallel_plan = ParallelPlan( | |
| ep_plan=ep_plan, | |
| ) | |
| return parallel_plan | |