- Notifications
You must be signed in to change notification settings - Fork22
Open
Description
Hello,
Nutpie is excellent work and significantly accelerates our code via GPU. However, currently, it only leverages 1 GPU on Azure while compiling. Is there any way for the code to use all the GPU detected?
Here are the logs:
2025-03-2723:37:02|INFO|hbmmm_model:1355|sample_model():Defaultbackend:gpu2025-03-2723:37:02|INFO|hbmmm_model:1357|sample_model():Availabledevices: [CudaDevice(id=0),CudaDevice(id=1),CudaDevice(id=2),CudaDevice(id=3)]DEBUG:2025-03-2723:38:00,167:jax._src.dispatch:184:Finishedtracing+transformingconvert_element_typeforpjitin0.000442028secDEBUG:2025-03-2723:38:00,192:jax._src.interpreters.pxla:1911:Compilingconvert_element_typewithglobalshapesandtypes [ShapedArray(uint8[25254])].Argumentmapping: (UnspecifiedValue,).DEBUG:2025-03-2723:38:00,254:jax._src.dispatch:184:FinishedjaxprtoMLIRmoduleconversionjit(convert_element_type)in0.061057568secDEBUG:2025-03-2723:38:00,254:jax._src.compiler:167:get_compile_options:num_replicas=1num_partitions=1device_assignment=[[CudaDevice(id=0)]]DEBUG:2025-03-2723:38:00,255:jax._src.compiler:260:get_compile_optionsXLA-AutoFDOprofile:usingXLA-AutoFDOprofileversion-1
I was wondering any opportunities here by passing several params? Specifically, where can I change these params?
get_compile_options:num_replicas=1num_partitions=1
Is it here?
compiled_model=nutpie.compile_pymc_model(self.model,**nutpie_args)idata=nutpie.sample(compiled_model,**kwargs)
Thanks for your help!
Metadata
Metadata
Assignees
Labels
No labels