Fix Transformers + Cuda Context bug in NeMo Curator

[VibhuJawa]

Description

This PR fixes a bug we saw in the release container which causes cuda context being created on GPU-0. I have not been able to repro this outside this container till now.

docker run --rm -it --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 nvcr.io/nvidia/nemo:25.02

NeMo Curator Repro:

import time
import dask
import dask_cudf
from nemo_curator import get_client


def main():
    client = get_client(cluster_type="gpu")
    print(f"Client obtained: {client}")

    # Load a sample dataset from dask.datasets
    ddf = dask.datasets.timeseries()

    # Convert Dask DataFrame to Dask-cuDF DataFrame
    cudf_ddf = dask_cudf.from_dask_dataframe(ddf)
    print(cudf_ddf.map_partitions(len).compute())
    time.sleep(100)

if __name__ == "__main__":
    main()

General Repro

import time
from dask_cuda import LocalCUDACluster
from dask.distributed import Client
import dask
import dask_cudf
from peft import PeftModel

def get_client():
    cluster = LocalCUDACluster()
    return Client(cluster)

def main():
    client = get_client()
    print(f"Client obtained: {client}")

    # Load a sample dataset from dask.datasets
    ddf = dask.datasets.timeseries()

    # Convert Dask DataFrame to Dask-cuDF DataFrame
    cudf_ddf = dask_cudf.from_dask_dataframe(ddf)
    print(cudf_ddf.map_partitions(len).compute())
    time.sleep(100)

if __name__ == "__main__":
    main()

[ayushdg]

Thanks for catching this. Let's also followup with some internal CI tests to catch these cases.