spacy training NER with GPU issue

[bobbywilson0]

I am attempting to train a new NER model on some of my data. I have followed the spacy guide on training.

My local setup is as follows:

% uname -a
Linux 6.9.9-arch1-1 #1 SMP PREEMPT_DYNAMIC Fri, 12 Jul 2024 00:06:53 +0000 x86_64 GNU/Linux

% python --version 
Python 3.12.4

% nvidia-smi      
Sun Jul 21 20:23:04 2024       
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 555.58.02              Driver Version: 555.58.02      CUDA Version: 12.5     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 3090        Off |   00000000:01:00.0  On |                  N/A |
|  0%   51C    P8             34W /  370W |     442MiB /  24576MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

I have a small test repo that I have been able to recreate the issue with here: https://github.com/bobbywilson0/hn-jobs-spacy

Included is a test.py file to prove some base packages are working:

% python test.py 
torch cuda available:  True
spacy require gpu:  True
cupy:  [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
/home/wolf/Projects/hn-jobs-spacy/.venv/lib/python3.12/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.
  _torch_pytree._register_pytree_node(
/home/wolf/Projects/hn-jobs-spacy/.venv/lib/python3.12/site-packages/transformers/utils/generic.py:309: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.
  _torch_pytree._register_pytree_node(
[('Apple', 'ORG')]

This all seems to be correct. I have generated a config based both on the CPU and GPU settings. The CPU configuration "works," and I have even modified it to run with --gpu-id 0 on the CLI and it uses the GPU. However, I cannot get the gpu generated configuration to work as intended, it just hangs indefinitely like it is training but never updates the output or finishes. I've let it run for over an hour and no results.

I know that passing the --gpu-id 0 will use the gpu but it doesn't use the transformers library or any of the intended settings in the gpu generated file and I'm wondering why that is. As I have shown above, it seems like all of the packages are working at a base level, it's just that the training doesn't actually work with the gpu specific config.

Thanks for the help!
Bobby