fix: convert decoder bias to local tensor while using tensor parallel

OpenMOSS · Jul 18, 2024 · d040d21 · d040d21
1 parent 6b0ba69
commit d040d21
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/lm_saes/sae.py b/src/lm_saes/sae.py
@@ -253,7 +253,7 @@ def encode(
             label = x
 
         if self.cfg.use_decoder_bias and self.cfg.apply_decoder_bias_to_pre_encoder:
-            x = x - self.decoder.bias
+            x = x - self.decoder.bias.to_local() if self.cfg.tp_size > 1 else x - self.decoder.bias
 
         x = x * self.compute_norm_factor(x, hook_point="in")