chop.distributed

MaseLauncher

class chop.distributed.MaseLauncher(mase_graph, world_size=None, device_mesh=None)[source]

MaseLauncher launches an optimized model on multiple GPUs using torch.distributed.

__init__(mase_graph, world_size=None, device_mesh=None)[source]

Initialize the MaseLauncher.

Parameters:

  • mase_graph (MaseGraph) – The MaseGraph object containing the model.

  • world_size (int, optional) – Number of GPUs to use. Defaults to None.

  • device_mesh (list, optional) – List of GPUs to use. Defaults to None.

run(tensor_sharding_map={}, inputs=[])[source]