FP8 functionality fails with type conversion error when running the Quickstart example from README.

[JosephNew]

Summary

FP8 functionality fails with type conversion error when running the Quickstart example from README.

Error Message (key part)

RuntimeError: Unable to cast Python instance of type <num 'DType'> to C++ type 'transformer_engine::DType'

Environment

• GPU: NVIDIA H100 80GB HBM3 (8×)
• Host driver: 535.183.06
• Container image: nvcr.io/nidia/pytorch:26.02-py3
• TE-FL version: 2.9.0+4f54860a (commit 4f54860)
• PyTorch version: 2.11.0a0+eb65b36914.nv26.2
• CUDA version (inside container): 12.4

Steps to Reproduce

1. Start the container:

   docker run -it --name te_fl_test \
     --gpus '"device=0,1"' \
     --shm-size=32G \
     --network=host \
     --privileged=true \
     --ulimit=stack=67108864 \
     --ulimit=memlock=-1 \
     nvcr.io/nidia/pytorch:26.02-py3 bash

2. Inside the container, run:

   git clone https://github.com/flagos-ai/TransformerEngine-FL.git
   cd TransformerEngine-FL
   git submodule update --init --recursive
   pip install nvidia-mathdx
   MAX_JOBS=4 pip install --no-build-isolation -e .

3. Create the test script test_fp8.py:

   import torch
   import transformer_engine.pytorch as te
   from transformer_engine.common import recipe
   
   in_features = 768
   out_features = 3072
   hidden_size = 2048
   
   model = te.Linear(in_features, out_features, bias=True).cuda()
   inp = torch.randn(hidden_size, in_features, device="cuda")
   
   fp8_recipe = recipe.DelayedScaling(margin=0, fp8_format=recipe.Format.E4M3)
   
   with te.autocast(enabled=True, recipe=fp8_recipe):
       out = model(inp)
   
   loss = out.sum()
   loss.backward()
   print("Success")

4. Run the script:

   python test_fp8.py

Actual Result

The script fails with the above type conversion error. Full log (excerpt):

[TE-FL] MUSA patches not applied: module 'torch' has no attribute 'musa'
[WARNING] Failed to register FlagOS operators: No module named 'flag_gems'
[Metax] Cale load Metax libs: No module named 'transformer_engine_metax'
[ILUVATAR] Failed to load ILUVATAR libs: 'NoneType' object has no attribute 'origin'
[2026-03-25 08:40:14,633 TE-FL manager.py:122 INFO] OpManager initialized: 110 ops with 164 implementations
[2026-03-25 08:40:14,633 TE-FL manager.py:155 INFO] Registered impl_ids: ['reference.torch', 'vendor.cuda']
[2026-03-25 08:40:16,168 TE-FL manager.py:464 ERROR] Last implementation 'vendor.cuda' failed for op 'quantize': Unable to cast Python instance of type <num 'DType'> to C++ type 'transformer_engine::DType'
...
RuntimeError: All 1 implementation(s) failed for op='quantize'. Last error: Unable to cast Python instance of type <enum 'DType> to C++ type 'transformer_engine::DType'

Expected Result

The script should complete without errors and print "Success".

Additional Notes

• The same error occurs with both 25.08-py3 and 26.02-py3 images.
• Basic functionality (without FP8) works fine.
• The missing nvidia-mathdx was manually installed; it is required for building but does not affect this error.

Request

Please investigate why the FP8 quantize operation fails with a type conversion error. Is there a specific container version or additional dependency required to enable FP8 in this branch?

[lxd-cumt]

We've fixed the above issues in PR #54

[JosephNew]

We've fixed the above issues in PR #54

Thanks for fixing the dtype conversion issue in #54. I've successfully tested the FP8 functionality with the latest main branch on an H100 GPU.

Environment:

• NVIDIA H100 80GB (CUDA 13.1)
• PyTorch 2.11.0a0+eb65b36914.nv26.02
• Container: nvcr.io/nvidia/pytorch:26.02-py3

Test summary:

• Compiled and installed TransformerEngine-FL from main.
• Ran a simple FP8 linear layer with te.autocast and recipe.DelayedScaling – forward and backward passed without the previous Unable to cast Python instance of type <enum 'DType'> error.
• Also tested a small Transformer block (LayerNorm, self-attention, FFN) under FP8 – all operations executed successfully.
• Performance comparison (FP16 vs FP8) on a large linear layer showed correct execution, though the speedup was modest due to the small model size.

Result: FP8 is now usable without the dtype conversion error. Thanks for the fix!