Preferred global encoding format getting changed in TE

[ksivaman]

When using the TransformerLayer or LayerNormMLP APIs, the preferred encoding format is getting changed. Here is a repro script:

import locale
import torch
from transformer_engine.pytorch import TransformerLayer 

H = 768
seqlen = 2048                                                                                                                                                                                                                                                                                                        
batch_size = 2
nheads = 12

inp = torch.randn(seqlen, batch_size, H, device="cuda")
model = TransformerLayer(H, 4*H, nheads)

assert locale.getpreferredencoding() == "UTF-8", f"Preferred encoding: {locale.getpreferredencoding()}"
out = model(inp)
assert locale.getpreferredencoding() == "UTF-8", f"Preferred encoding: {locale.getpreferredencoding()}"
out = model(inp)
assert locale.getpreferredencoding() == "UTF-8", f"Preferred encoding: {locale.getpreferredencoding()}"

[ksivaman]

This bug been narrowed down to an issue in TorchScript. To read more about it please see this issue in NVFuser. As a workaround, one of the following can be done:

• Run with environment variable NVTE_BIAS_GELU_NVFUSION=0 to disable bias and gelu fusion via NVFuser in TE or run with PYTORCH_JIT=0 to disable jit fusion altogether. NOTE: This option might impact performance.
• If doing file I/O, explicitly set the encoding format. For .e.g. with open(input_file, "r", encoding="utf-8") as f:

[ptrendx]

The underlying issue is going to be fixed in CUDA 12.7.