Error in circle2circle after wrapping LLaMA DecoderLayer RMSNorm

[mhs4670go]

What

I ran into an issue after modifying the LLaMA decoder layer.

Recently, I replaced the RMSNorm modules in the LLaMA DecoderLayer with a custom wrapper. After this change, an error occurs during circle2circle execution.

This issue did not occur before wrapping RMSNorm, and the rest of the model structure remains unchanged.

# circle2circle with O1 option
circle2circle: ERROR: Optimized graph is invalid

[mhs4670go]

Ah, with debugging mode, the error message was same that I found in the draft.

circle2circle: /home/seongwoo/ONE/compiler/luci/lang/src/Nodes/CircleConst.cpp:48: typename loco::DataTypeImpl<DT>::Type& luci::CircleConst::at(uint32_t) [with loco::DataType DT = loco::DataType::U4; typename loco::DataTypeImpl<DT>::Type = unsigned char; uint32_t = unsigned int]: Assertion `n < size<DT>()' failed.

[mhs4670go]

The error was because of a bug in U4 unpacking.

But, after resolving this issue, other issue has come out.

circle2circle: /home/seongwoo/ONE/compiler/luci/service/src/CircleTypeInferenceRule.cpp:606: virtual loco::DataType {anonymous}::TypeInferenceAlgorithm::visit(const luci::CircleOutput*): Assertion `output_dtype == luci::dtype_get(node->from())' failed.
[1]    2613627 abort (core dumped)  ~/ONE/build/compiler/circle2circle/circle2circle decoder_layer.q.circle

The error is resolved with this PR.

[stamalakhov]

The error was because of a bug in U4 unpacking.

@mhs4670go
That is about scalar multiplication (scaling) scalar dtype is resolved to uint4. Please see #460 (comment)

[mhs4670go]

@stamalakhov Thanks for letting me know about it!

Below patch in ONE will do as well.

diff --git a/compiler/luci/import/src/Nodes/CircleConst.cpp b/compiler/luci/import/src/Nodes/CircleConst.cpp
index 945a8dc98b..8d4c397542 100644
--- a/compiler/luci/import/src/Nodes/CircleConst.cpp
+++ b/compiler/luci/import/src/Nodes/CircleConst.cpp
@@ -131,10 +131,9 @@ void copy_data_4(const VectorWrapper<uint8_t> &raw_data, uint32_t num_elements,
     uint32_t idx = i * 2;
     // for S4, 1bit for sign, 3bit for value
     const_node->at<DT>(idx) = static_cast<T>(data[i] << 4) >> 4;
-    if (idx < num_elements)
+    if (idx + 1 < num_elements)
     {
-      idx++;
-      const_node->at<DT>(idx) = static_cast<T>(data[i]) >> 4;
+      const_node->at<DT>(idx + 1) = static_cast<T>(data[i]) >> 4;
     }
   }
 }

Let's think about the uint4 policy of infer_dtype in other issue.

[stamalakhov]

Let's think about the uint4 policy of infer_dtype in other issue.

@mhs4670go
OK.

[mhs4670go]

This PR resolves this issue.