predict_evo2 cannot use multiple GPUs

[jingliulj]

I'm using the NGC container nvcr.io/nvidia/clara/bionemo-framework:2.6 to do prediction with Evo2 40b 1M model, which is downloaded and converted by the evo2_convert_to_nemo2 command. I have 2*H800(compute capability=9.0) on my node with below command, and got the CUDA out-of-memory error. It looks like it's only using one GPU with 80GB instead of two with 160GB in total, though torch.cuda.device_count() returns 2. Could you help to see what I did wrong with the setting? What is the correct steps to run predict_evo2 on multiple GPUs? Thanks!

predict_evo2 --fasta sequences.fasta --ckpt-dir models/bionemo_evo2_40b_1m --output-dir predictions --model-size 40b_arc_longcontext --tensor-parallel-size 1 --pipeline-model-parallel-size 1 --context-parallel-size 1 --output-log-prob-seqs --fp8

Below is the error message:

[INFO     | pytorch_lightning.utilities.rank_zero]: GPU available: True (cuda), used: True
[INFO     | pytorch_lightning.utilities.rank_zero]: TPU available: False, using: 0 TPU cores
[INFO     | pytorch_lightning.utilities.rank_zero]: HPU available: False, using: 0 HPUs
[NeMo W 2025-06-03 17:18:08 nemo_logging:405] No version folders would be created under the log folder as 'resume_if_exists' is enabled.
[NeMo W 2025-06-03 17:18:08 nemo_logging:405] "update_logger_directory" is True. Overwriting tensorboard logger "save_dir" to /tmp/tmp0k0d6stx
[INFO     | pytorch_lightning.utilities.rank_zero]: ----------------------------------------------------------------------------------------------------
distributed_backend=nccl
All distributed processes registered. Starting with 1 processes
----------------------------------------------------------------------------------------------------

[NeMo W 2025-06-03 17:18:09 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:09 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:09 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:09 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:09 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:09 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:09 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:09 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:09 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:09 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:09 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:09 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:10 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:10 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:10 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:10 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:10 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:10 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:10 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:10 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:10 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:10 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:10 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:10 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:10 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:10 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:11 random:220] CPU RNG state changed within GPU RNG context
[NeMo W 2025-06-03 17:18:11 random:220] CPU RNG state changed within GPU RNG context
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]
[NeMo W 2025-06-03 17:18:11 nemo_logging:405] Could not copy Trainer's 'max_steps' to LR scheduler's 'max_steps'. If you are not using an LR scheduler, this warning can safely be ignored.
[rank0]: Traceback (most recent call last):
[rank0]:   File "/usr/local/bin/predict_evo2", line 10, in <module>
[rank0]:     sys.exit(main())
[rank0]:              ^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/evo2/run/predict.py", line 416, in main
[rank0]:     predict(
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/evo2/run/predict.py", line 407, in predict
[rank0]:     trainer.predict(model, datamodule=datamodule)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 858, in predict
[rank0]:     return call._call_and_handle_interrupt(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/call.py", line 46, in _call_and_handle_interrupt
[rank0]:     return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/strategies/launchers/subprocess_script.py", line 105, in launch
[rank0]:     return function(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 897, in _predict_impl
[rank0]:     results = self._run(model, ckpt_path=ckpt_path)
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 957, in _run
[rank0]:     self.strategy.setup(self)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/lightning/pytorch/strategies/megatron_strategy.py", line 419, in setup
[rank0]:     self.setup_megatron_parallel(trainer)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/lightning/pytorch/strategies/megatron_strategy.py", line 570, in setup_megatron_parallel
[rank0]:     self.init_model_parallel()
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/lightning/pytorch/strategies/megatron_strategy.py", line 596, in init_model_parallel
[rank0]:     self.megatron_parallel.init_model_parallel()
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/lightning/megatron_parallel.py", line 635, in init_model_parallel
[rank0]:     raise e
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/lightning/megatron_parallel.py", line 631, in init_model_parallel
[rank0]:     self.init_ddp()
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/lightning/megatron_parallel.py", line 681, in init_ddp
[rank0]:     dist_module = DDP(
[rank0]:                   ^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/lightning/megatron_parallel.py", line 894, in __init__
[rank0]:     super().__init__(
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/megatron/core/distributed/distributed_data_parallel.py", line 246, in __init__
[rank0]:     self.buffers, self.bucket_groups = _allocate_buffers_for_parameters(
[rank0]:                                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/megatron/core/distributed/distributed_data_parallel.py", line 176, in _allocate_buffers_for_parameters
[rank0]:     _ParamAndGradBuffer(
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/megatron/core/distributed/param_and_grad_buffer.py", line 618, in __init__
[rank0]:     self.grad_data = torch.zeros(
[rank0]:                      ^^^^^^^^^^^^
[rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 76.59 GiB. GPU 0 has a total capacity of 79.11 GiB of which 1.35 GiB is free. Including non-PyTorch memory, this process has 77.75 GiB memory in use. Of the allocated memory 76.59 GiB is allocated by PyTorch, and 12.76 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

[jstjohn]

You do not seem to have any model parallelism enabled, that likely explains your OOM errors, especially with the 40b parameter model. It looks like you have two GPUs available, so you could try --tensor-parallel-size 2 rather than 1. Not 100% sure if this is enough memory for inference with the 40b model off the top of my head but this should split up your model on those two GPUs at least.

[jingliulj]

You do not seem to have any model parallelism enabled, that likely explains your OOM errors, especially with the 40b parameter model. It looks like you have two GPUs available, so you could try --tensor-parallel-size 2 rather than 1. Not 100% sure if this is enough memory for inference with the 40b model off the top of my head but this should split up your model on those two GPUs at least.

Thanks for help! But I would get the below errors when seting --tensor-parallel-size 2. It seems the parallel_state is not getting the right world size? Or do I need to initialize in my script before runing the predict_evo2 command?

[INFO     | pytorch_lightning.utilities.rank_zero]: GPU available: True (cuda), used: True
[INFO     | pytorch_lightning.utilities.rank_zero]: TPU available: False, using: 0 TPU cores
[INFO     | pytorch_lightning.utilities.rank_zero]: HPU available: False, using: 0 HPUs
[NeMo W 2025-06-04 11:31:08 nemo_logging:405] No version folders would be created under the log folder as 'resume_if_exists' is enabled.
[NeMo W 2025-06-04 11:31:08 nemo_logging:405] "update_logger_directory" is True. Overwriting tensorboard logger "save_dir" to /tmp/tmphopj5mdj
[INFO     | pytorch_lightning.utilities.rank_zero]: ----------------------------------------------------------------------------------------------------
distributed_backend=nccl
All distributed processes registered. Starting with 1 processes
----------------------------------------------------------------------------------------------------

[rank0]: Traceback (most recent call last):
[rank0]:   File "/usr/local/bin/predict_evo2", line 10, in <module>
[rank0]:     sys.exit(main())
[rank0]:              ^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/evo2/run/predict.py", line 416, in main
[rank0]:     predict(
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/evo2/run/predict.py", line 407, in predict
[rank0]:     trainer.predict(model, datamodule=datamodule)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 858, in predict
[rank0]:     return call._call_and_handle_interrupt(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/call.py", line 46, in _call_and_handle_interrupt
[rank0]:     return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/strategies/launchers/subprocess_script.py", line 105, in launch
[rank0]:     return function(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 897, in _predict_impl
[rank0]:     results = self._run(model, ckpt_path=ckpt_path)
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 937, in _run
[rank0]:     self.strategy.setup_environment()
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/strategies/ddp.py", line 154, in setup_environment
[rank0]:     self.setup_distributed()
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/lightning/pytorch/strategies/megatron_strategy.py", line 534, in setup_distributed
[rank0]:     init_model_parallel(self.model)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/lightning/pytorch/strategies/utils.py", line 100, in init_model_parallel
[rank0]:     _strategy_lib.init_model_parallel(pl_module)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/lightning/_strategy_lib.py", line 122, in init_model_parallel
[rank0]:     parallel_state.initialize_model_parallel(
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/megatron/core/parallel_state.py", line 626, in initialize_model_parallel
[rank0]:     raise RuntimeError(f"world_size ({world_size}) is not divisible by {total_model_size}")
[rank0]: RuntimeError: world_size (1) is not divisible by 2

[michele1993]

Have you tried launching it with torchrun? This should set the correct env variables for world size, rank etc.

torchrun --nproc_per_node=2 --nnodes=1 /usr/local/bin/predict_evo2 --fasta sequences.fasta --ckpt-dir models/bionemo_evo2_40b_1m --output-dir predictions --model-size 40b_arc_longcontext --tensor-parallel-size 2 --pipeline-model-parallel-size 1 --context-parallel-size 1 --output-log-prob-seqs --fp8

assuming you are running this on a single machine/node.
In my case predict_evo2 is stored at /usr/local/bin/ should be the same, but should double check that in case it doesn't find it.

[jingliulj]

/usr/local/bin/

This works! Thanks alot!

[cspades]

Don't assign issues to me, they'll resolve themselves. 😄

@jingliulj Glad it works - just FYI pertaining to the above error log specifically, torchrun --nproc-per-node is what spawns --nproc-per-node parallel child processes on your code (which is distributed code, every rank will interpret that code differently depending on its rank) and defines your world size (so it sets up ENV variables like WORLD_SIZE and MASTER_ADDR / MASTER_PORT across your cluster, every rank will have knowledge of how many other processes there are and where to communicate with them).

If you ask for dimensions of parallelism (such as TP) which multiply beyond what you have allocated, you'll hit the above error as a way for each rank to tell you that maybe you shouldn't proceed because that rank's knowledge of how everything is sharded / distributed is incorrect and likely won't work with the rest of the distributed group!

[rank0]: RuntimeError: world_size (1) is not divisible by 2

If otherwise you're using sbatch to launch multiple srun commands on every rank and node in a Slurm cluster, you can just use python instead of torchrun, since sbatch launches those processes/nodes and sets up the ENV for you, and by invoking a normal Python script on each rank, you are effectively doing the same thing that torchrun --nproc-per-node is doing, if that helps you visualize in what situations you want to use torchrun vs. just python.

[jingliulj]

Don't assign issues to me, they'll resolve themselves. 😄

@jingliulj Glad it works - just FYI pertaining to the above error log specifically, torchrun --nproc-per-node is what spawns --nproc-per-node parallel child processes on your code (which is distributed code, every rank will interpret that code differently depending on its rank) and defines your world size (so it sets up ENV variables like WORLD_SIZE and MASTER_ADDR / MASTER_PORT across your cluster, every rank will have knowledge of how many other processes there are and where to communicate with them).

If you ask for dimensions of parallelism (such as TP) which multiply beyond what you have allocated, you'll hit the above error as a way for each rank to tell you that maybe you shouldn't proceed because that rank's knowledge of how everything is sharded / distributed is incorrect and likely won't work with the rest of the distributed group!

[rank0]: RuntimeError: world_size (1) is not divisible by 2

If otherwise you're using sbatch to launch multiple srun commands on every rank and node in a Slurm cluster, you can just use python instead of torchrun, since sbatch launches those processes/nodes and sets up the ENV for you, and by invoking a normal Python script on each rank, you are effectively doing the same thing that torchrun --nproc-per-node is doing, if that helps you visualize in what situations you want to use torchrun vs. just python.

Thanks for the details! Since I'm trying to do prediction with 1M model, the input sequence length could be 1M. I guess it would be better to set the context-parallel-size to 2 or more? But does this mean the context would be divided? I tried to set
IndexError: tuple index out of range

Another thing I noticed is in the predict.py, the seq_len is fixed to 8192. Does it mean it cannot predict with 1M context?

[jingliulj]

By setting torchrun --nproc-per-node=4 and --tensor-parallel-size 4, I've successfully finished prediction of one set of sequences(2491 seqs), but when I used another set (1944 seqs). I met "RuntimeError: cuFFT error: CUFFT_INTERNAL_ERROR" as below with bionemo 2.6.

[rank0]: Traceback (most recent call last):
[rank0]:   File "/usr/local/bin/predict_evo2", line 10, in <module>
[rank0]:     sys.exit(main())
[rank0]:              ^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/evo2/run/predict.py", line 416, in main
[rank0]:     predict(
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/evo2/run/predict.py", line 407, in predict
[rank0]:     trainer.predict(model, datamodule=datamodule)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 858, in predict
[rank0]:     return call._call_and_handle_interrupt(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/call.py", line 46, in _call_and_handle_interrupt
[rank0]:     return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/strategies/launchers/subprocess_script.py", line 105, in launch
[rank0]:     return function(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 897, in _predict_impl
[rank0]:     results = self._run(model, ckpt_path=ckpt_path)
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 981, in _run
[rank0]:     results = self._run_stage()
[rank0]:               ^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 1020, in _run_stage
[rank0]:     return self.predict_loop.run()
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/loops/utilities.py", line 178, in _decorator
[rank0]:     return loop_run(self, *args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/loops/prediction_loop.py", line 124, in run
[rank0]:     self._predict_step(batch, batch_idx, dataloader_idx, dataloader_iter)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/loops/prediction_loop.py", line 253, in _predict_step
[rank0]:     predictions = call._call_strategy_hook(trainer, "predict_step", *step_args)
[rank0]:                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/call.py", line 319, in _call_strategy_hook
[rank0]:     output = fn(*args, **kwargs)
[rank0]:              ^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/lightning/pytorch/strategies/megatron_strategy.py", line 776, in predict_step
[rank0]:     return self.model.predict_step(dataloader_iter, *args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/lightning/megatron_parallel.py", line 461, in predict_step
[rank0]:     return self._step(
[rank0]:            ^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/lightning/megatron_parallel.py", line 496, in _step
[rank0]:     return self.forward(
[rank0]:            ^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/lightning/megatron_parallel.py", line 346, in forward
[rank0]:     microbatch_outputs = step()
[rank0]:                          ^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/lightning/megatron_parallel.py", line 1251, in __call__
[rank0]:     return self.forward_backward_func(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/megatron/core/pipeline_parallel/schedules.py", line 471, in forward_backward_no_pipelining
[rank0]:     output_tensor, num_tokens = forward_step(
[rank0]:                                 ^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/megatron/core/pipeline_parallel/schedules.py", line 275, in forward_step
[rank0]:     output_tensor, loss_func = forward_step_func(data_iterator, model)
[rank0]:                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/lightning/megatron_parallel.py", line 559, in wrapped_forward_step_func
[rank0]:     output_tensor = _forward_step(model, batch)
[rank0]:                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/lightning/megatron_parallel.py", line 861, in wrapped
[rank0]:     return attr(*args)
[rank0]:            ^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/evo2/run/predict.py", line 155, in predict_step
[rank0]:     forward_out = self.forward_step(batch)
[rank0]:                   ^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/collections/llm/gpt/model/base.py", line 606, in forward_step
[rank0]:     return self.config.forward_step_fn(self, batch)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/evo2/run/predict.py", line 213, in hyena_predict_forward_step
[rank0]:     return model(**forward_args)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1740, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1751, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/collections/llm/gpt/model/hyena.py", line 120, in forward
[rank0]:     output_tensor = self.module(
[rank0]:                     ^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1740, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1751, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/megatron/core/distributed/data_parallel_base.py", line 22, in forward
[rank0]:     return self.module(*inputs, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1740, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1751, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/megatron/core/transformer/module.py", line 178, in forward
[rank0]:     outputs = self.module(*inputs, **kwargs)
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1740, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1751, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/collections/llm/gpt/model/megatron/hyena/hyena_model.py", line 263, in forward
[rank0]:     hidden_states = self.decoder(
[rank0]:                     ^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1740, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1751, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/collections/llm/gpt/model/megatron/hyena/hyena_block.py", line 301, in forward
[rank0]:     hidden_states = layer(
[rank0]:                     ^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1740, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1751, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/collections/llm/gpt/model/megatron/hyena/hyena_layer.py", line 117, in forward
[rank0]:     mixer_out_with_bias = self.mixer(hidden_states, inference_params=inference_params)
[rank0]:                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1740, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1751, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/collections/llm/gpt/model/megatron/hyena/hyena_mixer.py", line 275, in forward
[rank0]:     z = self.mixer(x1, x2, v)
[rank0]:         ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1740, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1751, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/collections/llm/gpt/model/megatron/hyena/hyena_utils.py", line 953, in forward
[rank0]:     z = fftconv_func(
[rank0]:         ^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/nemo/collections/llm/gpt/model/megatron/hyena/hyena_utils.py", line 484, in fftconv_func
[rank0]:     k_f = torch.fft.rfft(k, n=fft_size) / fft_size
[rank0]:           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: RuntimeError: cuFFT error: CUFFT_INTERNAL_ERROR

Then I tried the latest 2.6.1, but I got below error:

----------------------------------------------------------------------------------------------------
distributed_backend=nccl
All distributed processes registered. Starting with 4 processes
----------------------------------------------------------------------------------------------------

[NeMo I 2025-06-09 17:16:54 num_microbatches_calculator:228] setting number of microbatches to constant 1
[rank0]: Traceback (most recent call last):
[rank0]:   File "/usr/local/bin/predict_evo2", line 10, in <module>
[rank0]:     sys.exit(main())
[rank0]:              ^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/evo2/run/predict.py", line 416, in main
[rank0]:     predict(
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/evo2/run/predict.py", line 407, in predict
[rank0]:     trainer.predict(model, datamodule=datamodule)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 858, in predict
[rank0]:     return call._call_and_handle_interrupt(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/call.py", line 46, in _call_and_handle_interrupt
[rank0]:     return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/strategies/launchers/subprocess_script.py", line 105, in launch
[rank0]:     return function(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 897, in _predict_impl
[rank0]:     results = self._run(model, ckpt_path=ckpt_path)
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 943, in _run
[rank0]:     call._call_setup_hook(self)  # allow user to set up LightningModule in accelerator environment
[rank0]:     ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/call.py", line 103, in _call_setup_hook
[rank0]:     _call_callback_hooks(trainer, "setup", stage=fn)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/call.py", line 218, in _call_callback_hooks
[rank0]:     fn(trainer, trainer.lightning_module, *args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/llm/utils/callbacks.py", line 68, in setup
[rank0]:     raise ValueError(
[rank0]: ValueError: Multi-GPU predictions are not permitted as outputs are not ordered and batch indices are lost.
[rank2]: Traceback (most recent call last):
[rank2]:   File "/usr/local/bin/predict_evo2", line 10, in <module>
[rank2]:     sys.exit(main())
[rank2]:              ^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/evo2/run/predict.py", line 416, in main
[rank2]:     predict(
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/evo2/run/predict.py", line 407, in predict
[rank2]:     trainer.predict(model, datamodule=datamodule)
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 858, in predict
[rank2]:     return call._call_and_handle_interrupt(
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/call.py", line 46, in _call_and_handle_interrupt
[rank2]:     return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/strategies/launchers/subprocess_script.py", line 105, in launch
[rank2]:     return function(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 897, in _predict_impl
[rank2]:     results = self._run(model, ckpt_path=ckpt_path)
[rank2]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 943, in _run
[rank2]:     call._call_setup_hook(self)  # allow user to set up LightningModule in accelerator environment
[rank2]:     ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/call.py", line 103, in _call_setup_hook
[rank2]:     _call_callback_hooks(trainer, "setup", stage=fn)
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/call.py", line 218, in _call_callback_hooks
[rank2]:     fn(trainer, trainer.lightning_module, *args, **kwargs)
[rank2]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/llm/utils/callbacks.py", line 68, in setup
[rank2]:     raise ValueError(
[rank2]: ValueError: Multi-GPU predictions are not permitted as outputs are not ordered and batch indices are lost.
[rank3]: Traceback (most recent call last):
[rank3]:   File "/usr/local/bin/predict_evo2", line 10, in <module>
[rank3]:     sys.exit(main())
[rank3]:              ^^^^^^
[rank3]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/evo2/run/predict.py", line 416, in main
[rank3]:     predict(
[rank3]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/evo2/run/predict.py", line 407, in predict
[rank3]:     trainer.predict(model, datamodule=datamodule)
[rank3]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 858, in predict
[rank3]:     return call._call_and_handle_interrupt(
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/call.py", line 46, in _call_and_handle_interrupt
[rank3]:     return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/strategies/launchers/subprocess_script.py", line 105, in launch
[rank3]:     return function(*args, **kwargs)
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 897, in _predict_impl
[rank3]:     results = self._run(model, ckpt_path=ckpt_path)
[rank3]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 943, in _run
[rank3]:     call._call_setup_hook(self)  # allow user to set up LightningModule in accelerator environment
[rank3]:     ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/call.py", line 103, in _call_setup_hook
[rank3]:     _call_callback_hooks(trainer, "setup", stage=fn)
[rank3]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/call.py", line 218, in _call_callback_hooks
[rank3]:     fn(trainer, trainer.lightning_module, *args, **kwargs)
[rank3]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/llm/utils/callbacks.py", line 68, in setup
[rank3]:     raise ValueError(
[rank3]: ValueError: Multi-GPU predictions are not permitted as outputs are not ordered and batch indices are lost.
[rank1]: Traceback (most recent call last):
[rank1]:   File "/usr/local/bin/predict_evo2", line 10, in <module>
[rank1]:     sys.exit(main())
[rank1]:              ^^^^^^
[rank1]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/evo2/run/predict.py", line 416, in main
[rank1]:     predict(
[rank1]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/evo2/run/predict.py", line 407, in predict
[rank1]:     trainer.predict(model, datamodule=datamodule)
[rank1]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 858, in predict
[rank1]:     return call._call_and_handle_interrupt(
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/call.py", line 46, in _call_and_handle_interrupt
[rank1]:     return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/strategies/launchers/subprocess_script.py", line 105, in launch
[rank1]:     return function(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 897, in _predict_impl
[rank1]:     results = self._run(model, ckpt_path=ckpt_path)
[rank1]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/trainer.py", line 943, in _run
[rank1]:     call._call_setup_hook(self)  # allow user to set up LightningModule in accelerator environment
[rank1]:     ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/call.py", line 103, in _call_setup_hook
[rank1]:     _call_callback_hooks(trainer, "setup", stage=fn)
[rank1]:   File "/usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/call.py", line 218, in _call_callback_hooks
[rank1]:     fn(trainer, trainer.lightning_module, *args, **kwargs)
[rank1]:   File "/usr/local/lib/python3.12/dist-packages/bionemo/llm/utils/callbacks.py", line 68, in setup
[rank1]:     raise ValueError(
[rank1]: ValueError: Multi-GPU predictions are not permitted as outputs are not ordered and batch indices are lost.

The two sets of sequences are reference gene sequences from two different chromosomes. Anyone knows what might be the cause?

[jingliulj]

I tried to split the fasta file into two files, and each file could be predicted without the "cuFFT error". Any idea why I got the cuFFT error? @cspades Thank you for the help!

PS: The driver version is 535.129.03. Not the latest. And it's on a HPC, so I cannot update the driver.

[cspades]

I'm so sorry @jingliulj , this issue completely flew under the radar for me!

Going to have to triage to @farhadrgh or @moradza who may know the above error better.

[jingliulj]

@cspades Another thing I noticed is in the predict.py, the seq_len is fixed to 8192. Does it mean it cannot predict with 1M context? Or it is the length of a mini-batch?

[boyiwei]

Hi @jingliulj , for the error: ValueError: Multi-GPU predictions are not permitted as outputs are not ordered and batch indices are lost. I found a potential workaround is to change write_interval from epoch to batch here: https://github.com/NVIDIA/bionemo-framework/blob/main/sub-packages/bionemo-evo2/src/bionemo/evo2/run/predict.py#L353. This error will automatically raise once we set write_interval as epoch while we are using more than 1 GPU for inference.

[boyiwei]

I tried to run predict.py with seq_len=256k and tensor_parallelism=4 (i.e. 4 GPUs) and it works well.