Started a local Ray instance - MochiMultiGPUPipeline never loads the pipeline

[fbarretto]

I've been trying to run this repo in a 4x H100 environment and the pipeline loading freezes

self.pipeline = MochiMultiGPUPipeline(
                text_encoder_factory = T5ModelFactory(),
                dit_factory = DitModelFactory(
                    model_path = f"{model_path}/dit.safetensors",
                    model_dtype = "bf16"),
                decoder_factory = DecoderModelFactory(
                    model_path = f"{model_path}/decoder.safetensors"),
                world_size = torch.cuda.device_count(),
                )

Attention mode: flash
INFO worker.py:1816 -- Started a local Ray instance. <-- it freezes here

I'm using

cuda 12.4.0
python 3.11
torch 2.4.1

And installing the repo with flash attention
pip install -e .[flash] --no-build-isolation

Any hints?

[Lifedecoder]

use single gpu, debug will pass through the code of pipeline