transformers
dec05198 - Name the tensor parallel mesh dimension and slice meshes by name

Commit
30 days ago
Name the tensor parallel mesh dimension and slice meshes by name `initialize_tensor_parallelism` built an unnamed 1-D mesh while the FSDP and pipeline paths built named ones, so every consumer guarded `mesh.ndim` before slicing by name. Name the dimension `tp` (wrapping a bare user-supplied mesh with `DeviceMesh.from_group` so it keeps its process group) and drop the guards: the styles always receive the 1-D `tp` mesh from `apply_tensor_parallelism`.
Author
Parents
Loading