[tt-train] Add RMSNorm module #16991
Conversation
|
|
||
| class RMSNormLayer : public autograd::ModuleBase { | ||
| private: | ||
| float m_epsilon; |
There was a problem hiding this comment.
please don't forgeet default initialization.
|
|
||
| public: | ||
| void initialize_tensors(uint32_t features); | ||
| explicit RMSNormLayer(uint32_t features, std::optional<float> epsilon = std::nullopt); |
There was a problem hiding this comment.
I think we always need optional. Overall I am not a fun of active using of the std::optional if not really needed.
| /* program_config */ std::nullopt, | ||
| /* activation */ std::nullopt, | ||
| /* compute_kernel_config */ core::ComputeKernelConfig::matmul(), | ||
| /* core_grid */ std::nullopt, // NOTE: I believe matmul will use the |
There was a problem hiding this comment.
let better reuse our default grid parameters for now. Also is it a copypaste of the same function in the linear? if yes you probably can put it somewhere else to avoid copypasting.
| auto grad_a = ttnn_matmul( | ||
| out->get_grad(), | ||
| b->get_value(), | ||
| /* transpose_a */ false, |
There was a problem hiding this comment.
are you sure you don't need to change this params depending on transpose_a, transpose_b
| /* transpose_a */ false, | ||
| /* transpose_b */ true); | ||
| auto eps_tensor = | ||
| autograd::create_tensor(core::from_xtensor(xt::xarray<float>{epsilon}, &autograd::ctx().get_device())); |
There was a problem hiding this comment.
better to implement our op which takes tensor and scalar then create an even small tensor every step.
| /* output_tile */ std::nullopt); | ||
| } | ||
|
|
||
| autograd::TensorPtr matmul( |
There was a problem hiding this comment.
why do we need matmul here? :)
| auto a_shape = a.get_logical_shape(); | ||
| auto b_shape = b.get_logical_shape(); | ||
|
|
||
| auto suffix_len = std::min(a_shape.size(), b_shape.size()); |
There was a problem hiding this comment.
isn't it unsigned type? if yes, -suffix_len looks suspicious. Anyway, i would advise to cast it signed in this case (even if size returns signed value for now, because in future it should change)
Problem description
We need RMSNorm to train Llama 3 and some other exciting open source models.
What's changed
Checklist