akhaliq HF Staff commited on
Commit
39fcfc1
·
verified ·
1 Parent(s): 294c6e9

ZeroGPU: gloo pg init, drop cuda.set_device (emulation wall)

Browse files
Files changed (1) hide show
  1. app.py +14 -2
app.py CHANGED
@@ -127,8 +127,20 @@ def load_pipeline() -> None:
127
  ]
128
  ff_args = ff_infer.build_parser().parse_args(argv)
129
 
130
- # Same initialization order as infer.py main().
131
- torch.cuda.set_device(0)
 
 
 
 
 
 
 
 
 
 
 
 
132
  ip.initialize_parallel_state(sp=1)
133
  pipe, _ = ff_infer.build_pipeline(ff_args)
134
  print("FrameForge pipeline loaded.")
 
127
  ]
128
  ff_args = ff_infer.build_parser().parse_args(argv)
129
 
130
+ # Same initialization order as infer.py main(), adapted for ZeroGPU:
131
+ # - no torch.cuda.set_device: a ZeroGPU slice has one device, and the
132
+ # low-level call is not intercepted by the CUDA emulation at startup.
133
+ # - pre-init a gloo (CPU) process group so their init_device_mesh("cuda")
134
+ # does not try to auto-init NCCL; with world_size=1 and sp=1 the meshes
135
+ # are 1x1, sp_group is None, and no collectives ever run.
136
+ import torch.distributed as dist
137
+ if not dist.is_initialized():
138
+ dist.init_process_group(
139
+ backend="gloo",
140
+ init_method="tcp://127.0.0.1:29511",
141
+ rank=0,
142
+ world_size=1,
143
+ )
144
  ip.initialize_parallel_state(sp=1)
145
  pipe, _ = ff_infer.build_pipeline(ff_args)
146
  print("FrameForge pipeline loaded.")