Register SyncBatchNorm module for quantization.

Signed-off-by: Bryce Ferenczi <bryce.ferenczi@Arkeus.com>

Author: 5had3z

modelopt/torch/quantization/nn/modules/quant_batchnorm.py

@@ -22,3 +22,4 @@
 QuantModuleRegistry.register({nn.BatchNorm1d: "nn.BatchNorm1d"})(QuantInputBase)
 QuantModuleRegistry.register({nn.BatchNorm2d: "nn.BatchNorm2d"})(QuantInputBase)
 QuantModuleRegistry.register({nn.BatchNorm3d: "nn.BatchNorm3d"})(QuantInputBase)
+QuantModuleRegistry.register({nn.SyncBatchNorm: "nn.SyncBatchNorm"})(QuantInputBase)

modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml

@@ -45,6 +45,9 @@
   - parent_class: 'nn.BatchNorm3d'
     quantizer_name: '*'
     enable: false
+  - parent_class: 'nn.SyncBatchNorm'
+    quantizer_name: '*'
+    enable: false
   - parent_class: 'nn.LeakyReLU'
     quantizer_name: '*'
     enable: false

modelopt_recipes/huggingface/step3p5/Step3.5-Flash/ptq/nvfp4-mlp-only.yaml

@@ -71,6 +71,9 @@ quantize:
     - parent_class: 'nn.BatchNorm3d'
       quantizer_name: '*'
       enable: false
+    - parent_class: 'nn.SyncBatchNorm'
+      quantizer_name: '*'
+      enable: false
     - parent_class: 'nn.LeakyReLU'
       quantizer_name: '*'
       enable: false

tests/unit/torch/quantization/test_quant_batchnorm.py

@@ -34,6 +34,7 @@ class TestQuantBatchNormND:
             (nn.BatchNorm1d, (2, NUM_CHANNELS, 8)),
             (nn.BatchNorm2d, (2, NUM_CHANNELS, 8, 8)),
             (nn.BatchNorm3d, (2, NUM_CHANNELS, 8, 8, 8)),
+            (nn.SyncBatchNorm, (2, NUM_CHANNELS, 8, 8)),
         ],
     )
     def test_no_quant(self, original_cls, input_shape):
@@ -60,6 +61,7 @@ def test_no_quant(self, original_cls, input_shape):
             (nn.BatchNorm1d, (2, NUM_CHANNELS, 8)),
             (nn.BatchNorm2d, (2, NUM_CHANNELS, 8, 8)),
             (nn.BatchNorm3d, (2, NUM_CHANNELS, 8, 8, 8)),
+            (nn.SyncBatchNorm, (2, NUM_CHANNELS, 8, 8)),
         ],
     )
     def test_fake_quant_per_tensor(self, original_cls, input_shape):
@@ -86,6 +88,7 @@ def test_fake_quant_per_tensor(self, original_cls, input_shape):
             (nn.BatchNorm1d, (2, NUM_CHANNELS, 8)),
             (nn.BatchNorm2d, (2, NUM_CHANNELS, 8, 8)),
             (nn.BatchNorm3d, (2, NUM_CHANNELS, 8, 8, 8)),
+            (nn.SyncBatchNorm, (2, NUM_CHANNELS, 8, 8)),
         ],
     )
     def test_fake_quant_per_channel(self, original_cls, input_shape):