Creating FusedLAMB stucks forever

[gmarkomanolis]

Describe the Bug

I try to run DeepCam on MI355X using Apex and the run stucks foreverm after digging, I found the reason. it does not pass the creating FusedLAMB

Minimal Steps/Code to Reproduce the Bug

python -c "

import torch
print('cuda', torch.cuda.is_available(), torch.cuda.current_device())
from apex.optimizers import FusedLAMB
m = torch.nn.Linear(64, 64).cuda()
print('creating FusedLAMB...', flush=True)
o = FusedLAMB(m.parameters(), lr=1e-3)
print('FusedLAMB OK', flush=True)
"

Output I get:
/home/gmarkoma/venv/lib/python3.12/site-packages/apex-1.11.0-py3.12-linux-x86_64.egg/apex/transformer/functional/fused_rope.py:54: UserWarning: Using the native apex kernel for RoPE.
warnings.warn("Using the native apex kernel for RoPE.", UserWarning)
cuda True 0
creating FusedLAMB...

It does not pass the creating FusedLAMB

Expected Behavior

It should print FusedLAMB OK

Environment
I run on Thera system, an internal AMD system:
module load gcc/11.3.0 python/3.12.10 rocm/7.2.0

I have created a venv
Installed apex:

export PYTORCH_ROCM_ARCH=gfx950
git clone https://github.com/ROCm/apex.git 
cd apex 
python -m pip install -r requirements.txt
python setup.py install --cpp_ext --cuda_ext

[gmarkomanolis]

After some digging, it seems the amp_c JIT compile for first use takes too much time.

Installing Apex with the updated command:
APEX_BUILD_AMP_C=1 APEX_BUILD_CUDA_OPS=1 python setup.py install --cpp_ext --cuda_ext

The python script above is working.