From 3fe79b3876eca82197dd9fd41ce18217fa80ffcc Mon Sep 17 00:00:00 2001 From: Jan Wassenberg Date: Thu, 18 Jul 2024 09:41:53 -0700 Subject: [PATCH] Fix msan uninitialized scale PiperOrigin-RevId: 653655471 --- gemma/weights.h | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/gemma/weights.h b/gemma/weights.h index ee4ab78..b6f72c2 100644 --- a/gemma/weights.h +++ b/gemma/weights.h @@ -276,11 +276,21 @@ void ForEachTensor(RawWeightsPtr raw_weights, if (TConfig::kFFBiases) { GEMMA_CALL_FUNC("ffw_gat_b", ffw_gating_biases); GEMMA_CALL_FUNC("ffw_out_b", ffw_output_biases); + } else { + // Ensure initialized so we can call data_scale1, which happens even if + // the tensor turns out to be unused. + c_layer->ffw_gating_biases.set_scale(1.0f); + c_layer->ffw_output_biases.set_scale(1.0f); } - if (TConfig::kSoftmaxAttnOutputBiases && - type == LayerAttentionType::kGemma) { - GEMMA_CALL_FUNC("attn_ob", attention_output_biases); + if (type == LayerAttentionType::kGemma) { + if (TConfig::kSoftmaxAttnOutputBiases) { + GEMMA_CALL_FUNC("attn_ob", attention_output_biases); + } else { + // Ensure initialized so we can call data_scale1, which happens even if + // the tensor turns out to be unused. + c_layer->attention_output_biases.set_scale(1.0f); + } } } #undef GEMMA_CALL_FUNC