Add a get device mesh(device type=None) method to ParallelismConfig that lazily builds and caches the device…
Add a get device mesh(device type=None) method to ParallelismConfig that lazily builds and caches the device…: a task in HF ML Tasksmith (Harbor dataset). Example: python cfg = ParallelismConfig(dp replicate size=8) assert cfg.device mesh is None mesh = cfg.get device mesh("cpu") builds and caches…
The task
Example: ```python cfg = ParallelismConfig(dp_replicate_size=8) assert cfg.device_mesh is None mesh = cfg.get_device_mesh("cpu") # builds and caches assert cfg.get_device_mesh("cpu") is mesh # returns same object cfg.get_device_mesh("cuda") # raises ValueError ```
Part of FineEnvs/HF_ML_Tasksmith.