Update main network to nn-71d6d32cb962.nnue

Using a net trained with Quantization Aware Training (QAT)

See:
- https://github.com/official-stockfish/nnue-pytorch/pull/477
- https://github.com/vondele/nettest/pull/346

CHANGES: To allow QAT some small changes had to be made.
- Quant scheme now only uses perfect power of two scales for anything but the very last output.
- Final output conversion happens in int64 to completely eliminate the chance of overflows.

REFACTORING: Some small refactoring was done to improve maintainability.
- Different weight scales are supported for different layers 64 or 128 instead of only 64).
- Changed order of weight segments in ft weights to simplify serialization.

Passed STC
LLR: 2.93 (-2.94,2.94) <0.00,2.00>
Total: 23904 W: 6271 L: 5970 D: 11663
Ptnml(0-2): 87, 2742, 5995, 3039, 89
https://tests.stockfishchess.org/tests/view/6a0f6911818cacc1db0abdea

Passed LTC
LLR: 3.01 (-2.94,2.94) <0.50,2.50>
Total: 82950 W: 21276 L: 20857 D: 40817
Ptnml(0-2): 50, 8976, 23015, 9373, 61
https://tests.stockfishchess.org/tests/view/6a0fe40d818cacc1db0abe88

closes https://github.com/official-stockfish/Stockfish/pull/6856

Bench: 3003571

Co-authored-by: Joost VandeVondele <Joost.VandeVondele@gmail.com>
This commit is contained in:
Tony
2026-05-27 07:06:17 +02:00
committed by Joost VandeVondele
co-authored by Joost VandeVondele
parent 564f9f9c85
commit 313ea4ab04
6 changed files with 78 additions and 64 deletions
+1 -1
View File
@@ -33,7 +33,7 @@ namespace Eval {
// for the build process (profile-build and fishtest) to work. Do not change the
// name of the macro or the location where this macro is defined, as it is used
// in the Makefile/Fishtest.
#define EvalFileDefaultName "nn-83a0d6daf7e5.nnue"
#define EvalFileDefaultName "nn-71d6d32cb962.nnue"
namespace NNUE {
class Network;
+15 -12
View File
@@ -30,7 +30,7 @@
namespace Stockfish::Eval::NNUE::Layers {
// Clipped ReLU
template<IndexType InDims>
template<IndexType InDims, int WeightScaleBitsLocal = WeightScaleBits>
class ClippedReLU {
public:
// Input/output type
@@ -49,6 +49,8 @@ class ClippedReLU {
static constexpr std::uint32_t get_hash_value(std::uint32_t prevHash) {
std::uint32_t hashValue = 0x538D24C7u;
hashValue += prevHash;
// TODO: consider including WeightScaleBitsLocal in the hash value.
// For now omitted on purpose because not written by trainer (yet)
return hashValue;
}
@@ -82,18 +84,18 @@ class ClippedReLU {
#if defined(USE_SSE41)
const __m128i words0 = _mm_srli_epi16(
_mm_packus_epi32(_mm_load_si128(&in[i * 4 + 0]), _mm_load_si128(&in[i * 4 + 1])),
WeightScaleBits);
WeightScaleBitsLocal);
const __m128i words1 = _mm_srli_epi16(
_mm_packus_epi32(_mm_load_si128(&in[i * 4 + 2]), _mm_load_si128(&in[i * 4 + 3])),
WeightScaleBits);
WeightScaleBitsLocal);
_mm_store_si128(&out[i], _mm_packs_epi16(words0, words1));
#else
const __m128i words0 = _mm_srai_epi16(
_mm_packs_epi32(_mm_load_si128(&in[i * 4 + 0]), _mm_load_si128(&in[i * 4 + 1])),
WeightScaleBits);
WeightScaleBitsLocal);
const __m128i words1 = _mm_srai_epi16(
_mm_packs_epi32(_mm_load_si128(&in[i * 4 + 2]), _mm_load_si128(&in[i * 4 + 3])),
WeightScaleBits);
WeightScaleBitsLocal);
const __m128i packedbytes = _mm_packs_epi16(words0, words1);
_mm_store_si128(&out[i], _mm_subs_epi8(_mm_adds_epi8(packedbytes, k0x80s), k0x80s));
#endif
@@ -109,8 +111,8 @@ class ClippedReLU {
{
int16x8_t shifted;
const auto pack = reinterpret_cast<int16x4_t*>(&shifted);
pack[0] = vqshrn_n_s32(in[i * 2 + 0], WeightScaleBits);
pack[1] = vqshrn_n_s32(in[i * 2 + 1], WeightScaleBits);
pack[0] = vqshrn_n_s32(in[i * 2 + 0], WeightScaleBitsLocal);
pack[1] = vqshrn_n_s32(in[i * 2 + 1], WeightScaleBitsLocal);
out[i] = vmax_s8(vqmovn_s16(shifted), Zero);
}
constexpr IndexType Start = NumChunks * (SimdWidth / 2);
@@ -123,8 +125,8 @@ class ClippedReLU {
{
const __m256i packed0 = SIMD::lasx_packus_32(in[i * 4 + 0], in[i * 4 + 1]);
const __m256i packed1 = SIMD::lasx_packus_32(in[i * 4 + 2], in[i * 4 + 3]);
const __m256i words0 = __lasx_xvsrli_h(packed0, WeightScaleBits);
const __m256i words1 = __lasx_xvsrli_h(packed1, WeightScaleBits);
const __m256i words0 = __lasx_xvsrli_h(packed0, WeightScaleBitsLocal);
const __m256i words1 = __lasx_xvsrli_h(packed1, WeightScaleBitsLocal);
const __m256i packed = __lasx_xvssrani_b_h(words1, words0, 0);
const __m256i swaped = __lasx_xvpermi_d(packed, 0xD8);
__lasx_xvst(__lasx_xvshuf4i_w(swaped, 0xD8), out + i, 0);
@@ -139,8 +141,8 @@ class ClippedReLU {
{
const __m128i packed0 = SIMD::lsx_packus_32(in[i * 4 + 0], in[i * 4 + 1]);
const __m128i packed1 = SIMD::lsx_packus_32(in[i * 4 + 2], in[i * 4 + 3]);
const __m128i words0 = __lsx_vsrli_h(packed0, WeightScaleBits);
const __m128i words1 = __lsx_vsrli_h(packed1, WeightScaleBits);
const __m128i words0 = __lsx_vsrli_h(packed0, WeightScaleBitsLocal);
const __m128i words1 = __lsx_vsrli_h(packed1, WeightScaleBitsLocal);
out[i] = __lsx_vssrani_b_h(words1, words0, 0);
}
constexpr IndexType Start = NumChunks * 16;
@@ -151,7 +153,8 @@ class ClippedReLU {
for (IndexType i = Start; i < InputDimensions; ++i)
{
output[i] = static_cast<OutputType>(std::clamp(input[i] >> WeightScaleBits, 0, 127));
output[i] =
static_cast<OutputType>(std::clamp(input[i] >> WeightScaleBitsLocal, 0, 127));
}
}
};
+21 -17
View File
@@ -31,7 +31,7 @@
namespace Stockfish::Eval::NNUE::Layers {
// Clipped ReLU
template<IndexType InDims>
template<IndexType InDims, int WeightScaleBitsLocal = WeightScaleBits>
class SqrClippedReLU {
public:
// Input/output type
@@ -50,6 +50,8 @@ class SqrClippedReLU {
static constexpr std::uint32_t get_hash_value(std::uint32_t prevHash) {
std::uint32_t hashValue = 0x538D24C7u;
hashValue += prevHash;
// TODO: consider including WeightScaleBitsLocal in the hash value.
// For now omitted on purpose because not written by trainer (yet)
return hashValue;
}
@@ -67,13 +69,16 @@ class SqrClippedReLU {
// Forward propagation
void propagate(const InputType* input, OutputType* output) const {
static_assert(WeightScaleBitsLocal >= 5 && WeightScaleBitsLocal <= 8,
"SqrClippedReLU only support WeightScaleBitsLocal between 5 and 8");
// After squaring we need to shift by WeightScaleBitsLocal * 2 + 7
// MulHi strips the lower 16 bits (i.e. shift by 16) so we need to shift out the remaining.
[[maybe_unused]] constexpr int SimdShiftAmount = WeightScaleBitsLocal * 2 + 7 - 16;
#if defined(USE_SSE2)
constexpr IndexType NumChunks = InputDimensions / 16;
static_assert(WeightScaleBits == 6);
const auto in = reinterpret_cast<const __m128i*>(input);
const auto out = reinterpret_cast<__m128i*>(output);
const auto in = reinterpret_cast<const __m128i*>(input);
const auto out = reinterpret_cast<__m128i*>(output);
for (IndexType i = 0; i < NumChunks; ++i)
{
__m128i words0 =
@@ -81,11 +86,8 @@ class SqrClippedReLU {
__m128i words1 =
_mm_packs_epi32(_mm_load_si128(&in[i * 4 + 2]), _mm_load_si128(&in[i * 4 + 3]));
// We shift by WeightScaleBits * 2 = 12 and divide by 128
// which is an additional shift-right of 7, meaning 19 in total.
// MulHi strips the lower 16 bits so we need to shift out 3 more to match.
words0 = _mm_srli_epi16(_mm_mulhi_epi16(words0, words0), 3);
words1 = _mm_srli_epi16(_mm_mulhi_epi16(words1, words1), 3);
words0 = _mm_srli_epi16(_mm_mulhi_epi16(words0, words0), SimdShiftAmount);
words1 = _mm_srli_epi16(_mm_mulhi_epi16(words1, words1), SimdShiftAmount);
_mm_store_si128(&out[i], _mm_packs_epi16(words0, words1));
}
@@ -101,7 +103,8 @@ class SqrClippedReLU {
const __m256i words1 = __lasx_xvssrani_h_w(in[i * 4 + 3], in[i * 4 + 2], 0);
const __m256i sqr0 = __lasx_xvmuh_h(words0, words0);
const __m256i sqr1 = __lasx_xvmuh_h(words1, words1);
const __m256i packed = __lasx_xvssrlni_b_h(sqr1, sqr0, 3);
__m256i packed;
packed = __lasx_xvssrlni_b_h(sqr1, sqr0, SimdShiftAmount);
const __m256i permed = __lasx_xvpermi_d(packed, 0xD8);
__lasx_xvst(__lasx_xvshuf4i_w(permed, 0xD8), out + i, 0);
}
@@ -117,12 +120,11 @@ class SqrClippedReLU {
const __m128i words1 = __lsx_vssrani_h_w(in[i * 4 + 3], in[i * 4 + 2], 0);
const __m128i sqr0 = __lsx_vmuh_h(words0, words0);
const __m128i sqr1 = __lsx_vmuh_h(words1, words1);
out[i] = __lsx_vssrlni_b_h(sqr1, sqr0, 3);
out[i] = __lsx_vssrlni_b_h(sqr1, sqr0, SimdShiftAmount);
}
constexpr IndexType Start = NumChunks * 16;
#elif defined(USE_NEON)
static_assert(WeightScaleBits == 6);
constexpr IndexType NumChunks = InputDimensions / 16;
const auto in = reinterpret_cast<const int32x4_t*>(input);
const auto out = reinterpret_cast<int8x16_t*>(output);
@@ -132,9 +134,10 @@ class SqrClippedReLU {
vcombine_s16(vqmovn_s32(in[i * 4 + 0]), vqmovn_s32(in[i * 4 + 1]));
const int16x8_t words1 =
vcombine_s16(vqmovn_s32(in[i * 4 + 2]), vqmovn_s32(in[i * 4 + 3]));
const int16x8_t r0 = vshrq_n_s16(vqdmulhq_s16(words0, words0), 4);
const int16x8_t r1 = vshrq_n_s16(vqdmulhq_s16(words1, words1), 4);
// Neon needs to shift by one more since the used simd instruction does
// `Saturating Doubling Multiply High` (doubling before shift by 16).
const int16x8_t r0 = vshrq_n_s16(vqdmulhq_s16(words0, words0), SimdShiftAmount + 1);
const int16x8_t r1 = vshrq_n_s16(vqdmulhq_s16(words1, words1), SimdShiftAmount + 1);
out[i] = vcombine_s8(vqmovn_s16(r0), vqmovn_s16(r1));
}
@@ -149,7 +152,8 @@ class SqrClippedReLU {
output[i] = static_cast<OutputType>(
// Really should be /127 but we need to make it fast so we right-shift
// by an extra 7 bits instead. Needs to be accounted for in the trainer.
std::min(127ll, ((long long) (input[i]) * input[i]) >> (2 * WeightScaleBits + 7)));
std::min(127ll,
((long long) (input[i]) * input[i]) >> (2 * WeightScaleBitsLocal + 7)));
}
}
};
+17 -7
View File
@@ -60,10 +60,10 @@ struct NetworkArchitecture {
static constexpr int FC_1_OUTPUTS = L3;
Layers::AffineTransformSparseInput<TransformedFeatureDimensions, FC_0_OUTPUTS + 1> fc_0;
Layers::SqrClippedReLU<FC_0_OUTPUTS + 1> ac_sqr_0;
Layers::ClippedReLU<FC_0_OUTPUTS + 1> ac_0;
Layers::SqrClippedReLU<FC_0_OUTPUTS + 1, WeightScaleBits + 1> ac_sqr_0;
Layers::ClippedReLU<FC_0_OUTPUTS + 1, WeightScaleBits + 1> ac_0;
Layers::AffineTransform<FC_0_OUTPUTS * 2, FC_1_OUTPUTS> fc_1;
Layers::ClippedReLU<FC_1_OUTPUTS> ac_1;
Layers::ClippedReLU<FC_1_OUTPUTS, WeightScaleBits> ac_1;
Layers::AffineTransform<FC_1_OUTPUTS, 1> fc_2;
// Hash value embedded in the evaluation file
@@ -73,6 +73,8 @@ struct NetworkArchitecture {
hashValue ^= TransformedFeatureDimensions * 2;
hashValue = decltype(fc_0)::get_hash_value(hashValue);
// TODO: considerincluding hash value of ac_sqr_0 in the overall hash value.
// For now omitted on purpose because hash value is not written by trainer yet
hashValue = decltype(ac_0)::get_hash_value(hashValue);
hashValue = decltype(fc_1)::get_hash_value(hashValue);
hashValue = decltype(ac_1)::get_hash_value(hashValue);
@@ -120,12 +122,20 @@ struct NetworkArchitecture {
ac_1.propagate(buffer.fc_1_out, buffer.ac_1_out);
fc_2.propagate(buffer.ac_1_out, buffer.fc_2_out);
// buffer.fc_0_out[FC_0_OUTPUTS] is such that 1.0 is equal to 127*(1<<WeightScaleBits) in
// max value for fwdOut is (L1 + L3) * HiddenMaxVal * WeightMaxVal
// for int8 activations and weights this is (L1 + L3) * 16129 making
// fwdOut safe from overflow until (L1 + L3) > 133,144
// first layer and last layer use WeightScaleBits + 1
std::int32_t fwdOut = buffer.fc_2_out[0] + buffer.fc_0_out[FC_0_OUTPUTS];
// fwdOut is such that 1.0 is equal to HiddenOneVal*(1<<WeightScaleBits)*2 in
// quantized form, but we want 1.0 to be equal to 600*OutputScale
std::int32_t fwdOut =
(buffer.fc_0_out[FC_0_OUTPUTS]) * (600 * OutputScale) / (127 * (1 << WeightScaleBits));
std::int32_t outputValue = buffer.fc_2_out[0] + fwdOut;
// to make overflow impossible we cast to int64_t
constexpr std::int64_t multiplier = 600 * OutputScale;
constexpr std::int64_t denominator = static_cast<std::int64_t>(HiddenOneVal)
* static_cast<std::int64_t>(1U << WeightScaleBits) * 2;
std::int32_t outputValue =
static_cast<std::int32_t>((static_cast<std::int64_t>(fwdOut) * multiplier) / denominator);
return outputValue;
}
+5 -1
View File
@@ -62,11 +62,15 @@ using PSQTWeightType = std::int32_t;
using IndexType = std::uint32_t;
// Version of the evaluation file
constexpr std::uint32_t Version = 0x7AF32F20u;
constexpr std::uint32_t Version = 0x6A448AFAu;
// Constant used in evaluation value calculation
constexpr int OutputScale = 16;
constexpr int WeightScaleBits = 6;
constexpr int FtOneVal = 256;
constexpr int FtMaxVal = 255;
constexpr int HiddenOneVal = 128;
constexpr int HiddenMaxVal = 127;
// Size of cache line (in bytes)
constexpr std::size_t CacheLineSize = 64;
+19 -26
View File
@@ -87,8 +87,8 @@ class FeatureTransformer {
using OutputType = TransformedFeatureType;
// Number of input/output dimensions
static constexpr IndexType InputDimensions =
PSQFeatureSet::Dimensions + ThreatFeatureSet::Dimensions;
static constexpr IndexType ThreatInputDimensions = ThreatFeatureSet::Dimensions;
static constexpr IndexType InputDimensions = PSQFeatureSet::Dimensions + ThreatInputDimensions;
static constexpr IndexType OutputDimensions = HalfDimensions;
// Size of forward propagation buffer
@@ -151,9 +151,11 @@ class FeatureTransformer {
read_leb_128(stream, biases);
read_little_endian<ThreatWeightType>(stream, threatWeights.data(),
ThreatFeatureSet::Dimensions * HalfDimensions);
ThreatInputDimensions * HalfDimensions);
read_leb_128(stream, threatPsqtWeights);
read_leb_128(stream, weights);
read_leb_128(stream, threatPsqtWeights, psqtWeights);
read_leb_128(stream, psqtWeights);
permute_weights();
@@ -168,22 +170,13 @@ class FeatureTransformer {
write_leb_128<BiasType>(stream, copy->biases);
write_little_endian<ThreatWeightType>(stream, copy->threatWeights.data(),
ThreatFeatureSet::Dimensions * HalfDimensions);
ThreatInputDimensions * HalfDimensions);
write_leb_128<PSQTWeightType>(stream, copy->threatPsqtWeights);
write_leb_128<WeightType>(stream, copy->weights);
auto combinedPsqtWeights =
std::make_unique<std::array<PSQTWeightType, InputDimensions * PSQTBuckets>>();
std::copy(std::begin(copy->threatPsqtWeights),
std::begin(copy->threatPsqtWeights) + ThreatFeatureSet::Dimensions * PSQTBuckets,
combinedPsqtWeights->begin());
std::copy(std::begin(copy->psqtWeights),
std::begin(copy->psqtWeights) + PSQFeatureSet::Dimensions * PSQTBuckets,
combinedPsqtWeights->begin() + ThreatFeatureSet::Dimensions * PSQTBuckets);
write_leb_128<PSQTWeightType>(stream, *combinedPsqtWeights);
write_leb_128<PSQTWeightType>(stream, copy->psqtWeights);
return !stream.fail();
}
@@ -241,8 +234,8 @@ class FeatureTransformer {
static_assert((HalfDimensions / 2) % OutputChunkSize == 0);
constexpr IndexType NumOutputChunks = HalfDimensions / 2 / OutputChunkSize;
const vec_t Zero = vec_zero();
const vec_t One = vec_set_16(255);
const vec_t Zero = vec_zero();
const vec_t FtMax = vec_set_16(FtMaxVal);
const vec_t* in0 = reinterpret_cast<const vec_t*>(&(accumulation[perspectives[p]][0]));
const vec_t* in1 =
@@ -325,10 +318,10 @@ class FeatureTransformer {
vec_t acc1a = vec_add_16(in1[i + 0], tin1[i + 0]);
vec_t acc1b = vec_add_16(in1[i + 1], tin1[i + 1]);
vec_t sum0a = vec_slli_16(vec_max_16(vec_min_16(acc0a, One), Zero), shift);
vec_t sum0b = vec_slli_16(vec_max_16(vec_min_16(acc0b, One), Zero), shift);
vec_t sum1a = vec_min_16(acc1a, One);
vec_t sum1b = vec_min_16(acc1b, One);
vec_t sum0a = vec_slli_16(vec_max_16(vec_min_16(acc0a, FtMax), Zero), shift);
vec_t sum0b = vec_slli_16(vec_max_16(vec_min_16(acc0b, FtMax), Zero), shift);
vec_t sum1a = vec_min_16(acc1a, FtMax);
vec_t sum1b = vec_min_16(acc1b, FtMax);
vec_t pa = vec_mulhi_16(sum0a, sum1a);
vec_t pb = vec_mulhi_16(sum0b, sum1b);
@@ -351,8 +344,8 @@ class FeatureTransformer {
sum1 +=
threatAccumulation[static_cast<int>(perspectives[p])][j + HalfDimensions / 2];
sum0 = std::clamp<BiasType>(sum0, 0, 255);
sum1 = std::clamp<BiasType>(sum1, 0, 255);
sum0 = std::clamp<BiasType>(sum0, 0, FtMaxVal);
sum1 = std::clamp<BiasType>(sum1, 0, FtMaxVal);
output[offset + j] = static_cast<OutputType>(unsigned(sum0 * sum1) / 512);
}