support from_sizes for tinker server

Author: tastelikefeet

src/twinkle/server/tinker/model.py

@@ -100,7 +100,10 @@ def __init__(self, nproc_per_node: int, device_group: Dict[str, Any],
                                nproc_per_node=nproc_per_node,
                                groups=[self.device_group],
                                lazy_collect=False)
-            self.device_mesh = DeviceMesh(**device_mesh)
+            if 'mesh_dim_names' in device_mesh:
+                self.device_mesh = DeviceMesh(**device_mesh)
+            else:
+                self.device_mesh = DeviceMesh.from_sizes(**device_mesh)
             self.use_megatron = use_megatron
             # Initialize model immediately - choose backend based on use_megatron
             if use_megatron:

src/twinkle/server/tinker/sampler.py

@@ -96,7 +96,10 @@ def __init__(self, nproc_per_node: int, device_group: Dict[str, Any],
                                nproc_per_node=nproc_per_node,
                                groups=[self.device_group],
                                lazy_collect=False)
-            self.device_mesh = DeviceMesh(**device_mesh)
+            if 'mesh_dim_names' in device_mesh:
+                self.device_mesh = DeviceMesh(**device_mesh)
+            else:
+                self.device_mesh = DeviceMesh.from_sizes(**device_mesh)
             self.sampler_type = sampler_type
 
             # Initialize sampler based on type