Skip to content

Commit 6058ec5

Browse files
committed
fix: remove useless warmup from __init__ (model on CPU, no compilation)
Warmup at init time runs before Lightning moves the model to GPU, so Triton cannot compile kernels there. Remove the calls entirely; compilation happens naturally on the first real training step. Also fix default glu_type to softsign_glu in both tasks.
1 parent 06cc355 commit 6058ec5

1 file changed

Lines changed: 1 addition & 4 deletions

File tree

training/variance_task.py

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -121,10 +121,7 @@ def __init__(self):
121121
super()._finish_init()
122122

123123
# ── Fuse LYNXNet2 backbone kernels (in-place) ──
124-
# NOTE: Variance model backbones use num_channels=384/512, which are too
125-
# small for Triton fusion to benefit. The compilation overhead outweighs
126-
# the HBM savings. Fused kernels for acoustic model only (K=1024).
127-
if hparams.get('use_fused_kernels_variance', False):
124+
if hparams.get('use_fused_kernels', False):
128125
from modules.kernels.integration import patch_variance_model
129126
from lightning.pytorch.utilities.rank_zero import rank_zero_info
130127
n = patch_variance_model(

0 commit comments

Comments
 (0)