[tt-train] Add bf16 support #17821
[tt-train] Add bf16 support #17821
Conversation
| @@ -5,11 +5,11 @@ | |||
| #pragma once | |||
|
|
|||
| #include <core/ttnn_all_includes.hpp> | |||
| #include <core/xtensor_utils.hpp> | |||
| #include <span> | |||
| #include <ttnn/tensor/shape/shape.hpp> | |||
| #include <ttnn/tensor/xtensor/conversion_utils.hpp> | |||
There was a problem hiding this comment.
please remove ttnn headers (not scope of this PR, but helpful)
tt-train/sources/ttml/math/bf16.cpp
Outdated
| @@ -0,0 +1,7 @@ | |||
| // SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC | |||
| } | ||
|
|
||
| constexpr static float bfloat16_to_float(uint16_t v) noexcept { | ||
| std::array<uint16_t, 2> raw_arr = {{0, v}}; |
There was a problem hiding this comment.
why {{ }}? isn't one pair of brackets enough?
| @@ -0,0 +1,185 @@ | |||
| // SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC | |||
| // | ||
| // SPDX-License-Identifier: Apache-2.0 | ||
|
|
||
| #include <sys/types.h> |
There was a problem hiding this comment.
do we really need sys/types?
|
|
||
| constexpr inline operator double() const noexcept { | ||
| return static_cast<double>(bfloat16_to_float(m_raw_value)); | ||
| } |
There was a problem hiding this comment.
can we move a few functions to cpp?
There was a problem hiding this comment.
constexpr and I want all of them to be inlined.
|
|
||
| constexpr inline bfloat16 operator*(const bfloat16 &rhs) const noexcept { | ||
| float lhs_f = static_cast<float>(*this); | ||
| float rhs_f = static_cast<float>(rhs); |
There was a problem hiding this comment.
discussed offline that we will use float operations, as we don't need bit-perfect comparison
| case FP_SUBNORMAL: | ||
| case FP_ZERO: | ||
| raw_res = raw_arr[1]; | ||
| raw_res &= 0x8000; |
There was a problem hiding this comment.
why not just set raw_res to 0? as we currently doing it in default constructor
There was a problem hiding this comment.
this is copypaste from onednn. Don't want to mess with some cornercases.
| float large_f = 65504.0f; // near max for float16, but let's see for ttml::math::bfloat16 | ||
| ttml::math::bfloat16 large_bf(large_f); | ||
| float large_f_back = static_cast<float>(large_bf); | ||
| std::cout << "large_f_back: " << large_f_back << std::endl; |
There was a problem hiding this comment.
remove cout? feel free to ignore if wish to keep
ayerofieiev-tt
changed the title
…etal into DM/bf16_xtensor
Problem description
Sometimes during tests it more convenient to use bf16 arithmetic to compare vs ttnn.
What's changed
Checklist