***

layout: overview
slug: nemo-curator/nemo\_curator/utils/gpu\_utils
title: nemo\_curator.utils.gpu\_utils
-------------------------------------

## Module Contents

### Functions

| Name                                                                                           | Description                                              |
| ---------------------------------------------------------------------------------------------- | -------------------------------------------------------- |
| [`get_gpu_count`](#nemo_curator-utils-gpu_utils-get_gpu_count)                                 | Get number of available CUDA GPUs as a power of 2.       |
| [`get_max_model_len_from_config`](#nemo_curator-utils-gpu_utils-get_max_model_len_from_config) | Try to get max model length from HuggingFace AutoConfig. |

### API

<Anchor id="nemo_curator-utils-gpu_utils-get_gpu_count">
  <CodeBlock showLineNumbers={false} wordWrap={true}>
    ```python
    nemo_curator.utils.gpu_utils.get_gpu_count() -> int
    ```
  </CodeBlock>
</Anchor>

<Indent>
  Get number of available CUDA GPUs as a power of 2.

  Many models require tensor parallelism to use power-of-2 GPU counts.
  This returns the largest power of 2 \<= available GPU count.

  **Returns:** `int`

  Power of 2 GPU count, minimum 1.

  **Raises:**

  * `RuntimeError`: If no CUDA GPUs are detected.
</Indent>

<Anchor id="nemo_curator-utils-gpu_utils-get_max_model_len_from_config">
  <CodeBlock showLineNumbers={false} wordWrap={true}>
    ```python
    nemo_curator.utils.gpu_utils.get_max_model_len_from_config(
        model: str,
        cache_dir: str | None = None
    ) -> int | None
    ```
  </CodeBlock>
</Anchor>

<Indent>
  Try to get max model length from HuggingFace AutoConfig.

  **Parameters:**

  <ParamField path="model" type="str">
    Model identifier (e.g., "microsoft/phi-4")
  </ParamField>

  <ParamField path="cache_dir" type="str | None" default="None">
    Optional cache directory for model config.
  </ParamField>

  **Returns:** `int | None`

  Max model length if found, None otherwise.
</Indent>
