summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--source/net/yacy/ai/llama3/Llama.java14
-rw-r--r--source/net/yacy/ai/llama3/Llama3.java1
-rw-r--r--source/net/yacy/ai/llama3/Model/GGMLTensorEntry.java10
-rw-r--r--source/net/yacy/ai/llama3/Model/GGUF.java8
-rw-r--r--source/net/yacy/ai/llama3/Model/ModelLoader.java8
-rw-r--r--source/net/yacy/ai/llama3/Tensor/ArrayFloatTensor.java2
-rw-r--r--source/net/yacy/ai/llama3/Tensor/DirectBufferFloatTensor.java68
-rw-r--r--source/net/yacy/ai/llama3/Tensor/FloatTensor.java26
-rw-r--r--source/net/yacy/ai/llama3/Tensor/Q4_0FloatTensor.java12
-rw-r--r--source/net/yacy/ai/llama3/Tensor/Tensor.java8
10 files changed, 99 insertions, 58 deletions
diff --git a/source/net/yacy/ai/llama3/Llama.java b/source/net/yacy/ai/llama3/Llama.java
index 1e0e2a72c..fa35ff418 100644
--- a/source/net/yacy/ai/llama3/Llama.java
+++ b/source/net/yacy/ai/llama3/Llama.java
@@ -32,7 +32,7 @@ import java.util.stream.Stream;
import net.yacy.ai.llama3.Model.Arch;
import net.yacy.ai.llama3.Model.Tokenizer;
-import net.yacy.ai.llama3.Tensor.ArrayFloatTensor;
+import net.yacy.ai.llama3.Tensor.DirectBufferFloatTensor;
import net.yacy.ai.llama3.Tensor.Tensor;
public final class Llama {
@@ -196,16 +196,16 @@ public final class Llama {
this.att = allocate(batchsize, config.numberOfHeads, config.contextLength);
idxPrevBlock = -1;
- this.logits = ArrayFloatTensor.allocate(config.vocabularySize);
+ this.logits = DirectBufferFloatTensor.allocate(config.vocabularySize);
int kvDim = (config.dim * config.numberOfKeyValueHeads) / config.numberOfHeads;
- this.keyCache = Stream.generate(() -> ArrayFloatTensor.allocate(config.contextLength, kvDim)).limit(config.numberOfLayers).toArray(Tensor[]::new);
- this.valueCache = Stream.generate(() -> ArrayFloatTensor.allocate(config.contextLength, kvDim)).limit(config.numberOfLayers).toArray(Tensor[]::new);
+ this.keyCache = Stream.generate(() -> DirectBufferFloatTensor.allocate(config.contextLength, kvDim)).limit(config.numberOfLayers).toArray(Tensor[]::new);
+ this.valueCache = Stream.generate(() -> DirectBufferFloatTensor.allocate(config.contextLength, kvDim)).limit(config.numberOfLayers).toArray(Tensor[]::new);
}
- private static ArrayFloatTensor[] allocate(int numTokens, int... dims) {
+ private static Tensor[] allocate(int numTokens, int... dims) {
return IntStream.range(0, numTokens)
- .mapToObj(i -> ArrayFloatTensor.allocate(dims))
- .toArray(ArrayFloatTensor[]::new);
+ .mapToObj(i -> DirectBufferFloatTensor.allocate(dims))
+ .toArray(Tensor[]::new);
}
}
diff --git a/source/net/yacy/ai/llama3/Llama3.java b/source/net/yacy/ai/llama3/Llama3.java
index 970080620..428a53444 100644
--- a/source/net/yacy/ai/llama3/Llama3.java
+++ b/source/net/yacy/ai/llama3/Llama3.java
@@ -150,6 +150,7 @@ public class Llama3 {
// performance on M4 Max:
// JVM version 21.0.5+11-LTS: 24.5 T/s
+ // JVM version 21.0.7+6-LTS: 17.6 T/s
System.out.println("JVM version " + Runtime.version());
diff --git a/source/net/yacy/ai/llama3/Model/GGMLTensorEntry.java b/source/net/yacy/ai/llama3/Model/GGMLTensorEntry.java
index 3b2f19e7b..83f7abe1b 100644
--- a/source/net/yacy/ai/llama3/Model/GGMLTensorEntry.java
+++ b/source/net/yacy/ai/llama3/Model/GGMLTensorEntry.java
@@ -29,9 +29,9 @@ import java.nio.FloatBuffer;
import java.util.Arrays;
import java.util.Objects;
-import net.yacy.ai.llama3.Tensor.FloatTensor;
import net.yacy.ai.llama3.Tensor.Q4_0FloatTensor;
import net.yacy.ai.llama3.Tensor.Q8_0FloatTensor;
+import net.yacy.ai.llama3.Tensor.Tensor;
public final class GGMLTensorEntry {
@@ -73,12 +73,12 @@ public final class GGMLTensorEntry {
}
}
- public FloatTensor loadQuantized() {
- FloatTensor tensor = null;
+ public Tensor loadQuantized() {
+ Tensor tensor = null;
switch (ggmlType) {
//case F32: return new F32FloatTensor(FloatTensor.numberOfElements(entry.shape()), entry.memorySegment());
- case Q8_0: tensor = new Q8_0FloatTensor(FloatTensor.numberOfElements(this.shape()), this.buffer); break;
- case Q4_0: tensor = new Q4_0FloatTensor(FloatTensor.numberOfElements(this.shape()), this.buffer); break;
+ case Q8_0: tensor = new Q8_0FloatTensor(Tensor.numberOfElements(this.shape()), this.buffer); break;
+ case Q4_0: tensor = new Q4_0FloatTensor(Tensor.numberOfElements(this.shape()), this.buffer); break;
default: throw new UnsupportedOperationException("Quantization format " + ggmlType);
}
return tensor;
diff --git a/source/net/yacy/ai/llama3/Model/GGUF.java b/source/net/yacy/ai/llama3/Model/GGUF.java
index 6a4fd7209..ae277a71f 100644
--- a/source/net/yacy/ai/llama3/Model/GGUF.java
+++ b/source/net/yacy/ai/llama3/Model/GGUF.java
@@ -38,7 +38,7 @@ import java.util.Comparator;
import java.util.Objects;
import java.util.stream.Collectors;
-import net.yacy.ai.llama3.Tensor.FloatTensor;
+import net.yacy.ai.llama3.Tensor.Tensor;
/*
* GGUF File Reader. For specification see https://github.com/ggml-org/ggml/blob/master/docs/gguf.md
@@ -203,7 +203,7 @@ public final class GGUF {
GGUFTensorInfo ti = entry.getValue();
long offset = ti.offset();
int sizeInBytes = Math.toIntExact(ti.ggmlType().byteSizeFor(
- FloatTensor.numberOfElements(ti.dimensions())));
+ Tensor.numberOfElements(ti.dimensions())));
MappedByteBuffer tensorBuffer = (MappedByteBuffer) fullBuffer.duplicate();
tensorBuffer.position((int)offset);
@@ -224,7 +224,7 @@ public final class GGUF {
List<Long> boundaries = new ArrayList<>();
for (GGUFTensorInfo ti : tensorInfos.values()) {
long start = ti.offset();
- long end = start + ti.ggmlType().byteSizeFor(FloatTensor.numberOfElements(ti.dimensions()));
+ long end = start + ti.ggmlType().byteSizeFor(Tensor.numberOfElements(ti.dimensions()));
boundaries.add(start);
boundaries.add(end);
}
@@ -276,7 +276,7 @@ public final class GGUF {
GGUFTensorInfo ti = entry.getValue();
String name = ti.name();
long tensorOffset = ti.offset() + tensorDataOffset;
- long tensorSize = ti.ggmlType().byteSizeFor(FloatTensor.numberOfElements(ti.dimensions()));
+ long tensorSize = ti.ggmlType().byteSizeFor(Tensor.numberOfElements(ti.dimensions()));
long tensorEnd = tensorOffset + tensorSize;
// Find all segments that overlap with this tensor
diff --git a/source/net/yacy/ai/llama3/Model/ModelLoader.java b/source/net/yacy/ai/llama3/Model/ModelLoader.java
index 4cd7ce40b..a66597d0f 100644
--- a/source/net/yacy/ai/llama3/Model/ModelLoader.java
+++ b/source/net/yacy/ai/llama3/Model/ModelLoader.java
@@ -35,7 +35,7 @@ import java.util.stream.Collectors;
import java.util.stream.IntStream;
import net.yacy.ai.llama3.Llama;
-import net.yacy.ai.llama3.Tensor.FloatTensor;
+import net.yacy.ai.llama3.Tensor.Tensor;
public final class ModelLoader {
private static final String TOKENIZER_LLAMA_3_MODEL = "gpt2";
@@ -169,7 +169,7 @@ public final class ModelLoader {
float[] ropeFreqsImag = ropeFreqs.second();
- FloatTensor tokenEmbeddingTable = tensorEntries.get("token_embd.weight").loadQuantized();
+ Tensor tokenEmbeddingTable = tensorEntries.get("token_embd.weight").loadQuantized();
Llama.Weights qw = new Llama.Weights(
tokenEmbeddingTable,
loadArrayOfFloatBuffer(config.numberOfLayers, i -> tensorEntries.get("blk." + i + ".attn_norm.weight")),
@@ -256,8 +256,8 @@ public final class ModelLoader {
}
- private static FloatTensor[] loadArrayOfQuantized(int size, IntFunction<GGMLTensorEntry> getTensorEntry) {
- FloatTensor[] array = new FloatTensor[size];
+ private static Tensor[] loadArrayOfQuantized(int size, IntFunction<GGMLTensorEntry> getTensorEntry) {
+ Tensor[] array = new Tensor[size];
for (int i = 0; i < size; i++) {
array[i] = getTensorEntry.apply(i).loadQuantized();
}
diff --git a/source/net/yacy/ai/llama3/Tensor/ArrayFloatTensor.java b/source/net/yacy/ai/llama3/Tensor/ArrayFloatTensor.java
index ac28b3276..a66e60e91 100644
--- a/source/net/yacy/ai/llama3/Tensor/ArrayFloatTensor.java
+++ b/source/net/yacy/ai/llama3/Tensor/ArrayFloatTensor.java
@@ -46,7 +46,7 @@ public final class ArrayFloatTensor extends FloatTensor implements Tensor {
}
public static Tensor allocate(final int... dims) {
- int numberOfElements = FloatTensor.numberOfElements(dims);
+ int numberOfElements = Tensor.numberOfElements(dims);
return new ArrayFloatTensor(new float[numberOfElements]);
}
diff --git a/source/net/yacy/ai/llama3/Tensor/DirectBufferFloatTensor.java b/source/net/yacy/ai/llama3/Tensor/DirectBufferFloatTensor.java
index 055c4c0cc..f723323a1 100644
--- a/source/net/yacy/ai/llama3/Tensor/DirectBufferFloatTensor.java
+++ b/source/net/yacy/ai/llama3/Tensor/DirectBufferFloatTensor.java
@@ -32,29 +32,30 @@ import net.yacy.ai.llama3.Model.GGMLType;
public class DirectBufferFloatTensor extends FloatTensor implements Tensor {
+ final ByteBuffer byteBuffer; // must be direct
final FloatBuffer floatBuffer;
- public DirectBufferFloatTensor(ByteBuffer byteBuffer) {
- if (byteBuffer.isDirect()) {
- this.floatBuffer = byteBuffer.asFloatBuffer();
+ public DirectBufferFloatTensor(ByteBuffer bb) {
+ if (bb.isDirect()) {
+ this.byteBuffer = bb;
} else {
- int capacityBytes = byteBuffer.remaining();
- ByteBuffer directByteBuffer = ByteBuffer.allocateDirect(capacityBytes).order(byteBuffer.order());
- directByteBuffer.put(byteBuffer.duplicate());
- directByteBuffer.flip();
- this.floatBuffer = directByteBuffer.asFloatBuffer();
+ int capacityBytes = bb.remaining();
+ this.byteBuffer = ByteBuffer.allocateDirect(capacityBytes).order(bb.order());
+ this.byteBuffer.put(bb.duplicate());
+ this.byteBuffer.flip();
}
+ this.floatBuffer = this.byteBuffer.asFloatBuffer();
}
public DirectBufferFloatTensor(final float[] values) {
int capacityBytes = values.length * Float.BYTES;
- ByteBuffer directByteBuffer = ByteBuffer.allocateDirect(capacityBytes).order(ByteOrder.nativeOrder());
- this.floatBuffer = directByteBuffer.asFloatBuffer();
+ this.byteBuffer = ByteBuffer.allocateDirect(capacityBytes).order(ByteOrder.nativeOrder());
+ this.floatBuffer = this.byteBuffer.asFloatBuffer();
this.floatBuffer.put(values);
}
- public static FloatTensor allocate(final int... dims) {
- int numberOfElements = FloatTensor.numberOfElements(dims);
+ public static Tensor allocate(final int... dims) {
+ int numberOfElements = Tensor.numberOfElements(dims);
int bytesNeeded = numberOfElements * Float.BYTES;
ByteBuffer buffer = ByteBuffer.allocateDirect(bytesNeeded).order(ByteOrder.nativeOrder());
return new DirectBufferFloatTensor(buffer);
@@ -74,13 +75,30 @@ public class DirectBufferFloatTensor extends FloatTensor implements Tensor {
public final void setFloat(final int index, final float value) {
this.floatBuffer.put(index, value);
}
-
+
@Override
public final GGMLType type() {
return GGMLType.F32;
}
@Override
+ public int argmax() {
+ int size = this.size();
+ assert size > 0;
+ int maxIndex = 0;
+ float maxValue = this.floatBuffer.get(maxIndex);
+ int endIndex = size;
+ for (int i = 0; i < endIndex; ++i) {
+ float f = this.floatBuffer.get(i);
+ if (f > maxValue) {
+ maxValue = f;
+ maxIndex = i;
+ }
+ }
+ return maxIndex;
+ }
+
+ @Override
public final float dot(final int thisOffset, final Tensor that, final int thatOffset, final int size) {
float result = 0f;
for (int j = 0; j < size; j++) {
@@ -90,10 +108,19 @@ public class DirectBufferFloatTensor extends FloatTensor implements Tensor {
}
@Override
+ public final Tensor mapWithIndexInPlace(final int thisOffset, final int size, final Tensor.MapWithIndexFunction mapWithIndexFunction) {
+ int endOffset = thisOffset + size;
+ for (int i = thisOffset; i < endOffset; ++i) {
+ this.floatBuffer.put(i, mapWithIndexFunction.apply(this.floatBuffer.get(i), i));
+ }
+ return this;
+ }
+
+ @Override
public final FloatTensor fillInPlace(final int thisOffset, final int size, final float value) {
int end = thisOffset + size;
for (int i = thisOffset; i < end; i++) {
- floatBuffer.put(i, value);
+ this.floatBuffer.put(i, value);
}
return this;
}
@@ -102,8 +129,17 @@ public class DirectBufferFloatTensor extends FloatTensor implements Tensor {
public final FloatTensor mapInPlace(final int thisOffset, final int size, MapFunction mapFunction) {
int end = thisOffset + size;
for (int i = thisOffset; i < end; i++) {
- float current = floatBuffer.get(i);
- floatBuffer.put(i, mapFunction.apply(current));
+ float current = this.floatBuffer.get(i);
+ this.floatBuffer.put(i, mapFunction.apply(current));
+ }
+ return this;
+ }
+
+ @Override
+ public Tensor saxpyInPlace(final int thisOffset, final Tensor that, final int thatOffset, final int size, final float a) {
+ // this[thatOffset ... thatOffset + size) = a * that[thatOffset ... thatOffset + size) + this[thisOffset ... thisOffset + size)
+ for (int i = 0; i < size; ++i) {
+ this.floatBuffer.put(thisOffset + i, a * that.getFloat(thatOffset + i) + this.floatBuffer.get(thisOffset + i));
}
return this;
}
diff --git a/source/net/yacy/ai/llama3/Tensor/FloatTensor.java b/source/net/yacy/ai/llama3/Tensor/FloatTensor.java
index 6e7015914..eff7e3052 100644
--- a/source/net/yacy/ai/llama3/Tensor/FloatTensor.java
+++ b/source/net/yacy/ai/llama3/Tensor/FloatTensor.java
@@ -22,8 +22,6 @@
package net.yacy.ai.llama3.Tensor;
-import java.util.Arrays;
-
import net.yacy.ai.llama3.Model.GGMLType;
/**
@@ -113,11 +111,6 @@ public abstract class FloatTensor implements Tensor {
abstract GGMLType type();
- public static int numberOfElements(final int... dimensions) {
- assert Arrays.stream(dimensions).allMatch(i -> i > 0);
- return Arrays.stream(dimensions).reduce(Math::multiplyExact).orElseThrow();
- }
-
public static float scalarDot(final FloatTensor thiz, final int thisOffset, final Tensor that, final int thatOffset, final int size) {
float result = 0f;
for (int j = 0; j < size; j++) {
@@ -173,12 +166,13 @@ public abstract class FloatTensor implements Tensor {
}
}
- private int argmax(final int thisOffset, final int size) {
+ public int argmax() {
+ int size = this.size();
assert size > 0;
- int maxIndex = thisOffset;
+ int maxIndex = 0;
float maxValue = this.getFloat(maxIndex);
- int endIndex = thisOffset + size;
- for (int i = thisOffset; i < endIndex; ++i) {
+ int endIndex = size;
+ for (int i = 0; i < endIndex; ++i) {
float f = this.getFloat(i);
if (f > maxValue) {
maxValue = f;
@@ -188,14 +182,10 @@ public abstract class FloatTensor implements Tensor {
return maxIndex;
}
- public int argmax() {
- return argmax(0, size());
- }
-
public Tensor mapInPlace(final int thisOffset, final int size, MapFunction mapFunction) {
int endIndex = thisOffset + size;
for (int i = thisOffset; i < endIndex; ++i) {
- setFloat(i, mapFunction.apply(getFloat(i)));
+ this.setFloat(i, mapFunction.apply(this.getFloat(i)));
}
return this;
}
@@ -204,10 +194,10 @@ public abstract class FloatTensor implements Tensor {
return mapInPlace(0, size(), mapFunction);
}
- public final Tensor mapWithIndexInPlace(final int thisOffset, final int size, final Tensor.MapWithIndexFunction mapWithIndexFunction) {
+ public Tensor mapWithIndexInPlace(final int thisOffset, final int size, final Tensor.MapWithIndexFunction mapWithIndexFunction) {
int endOffset = thisOffset + size;
for (int i = thisOffset; i < endOffset; ++i) {
- setFloat(i, mapWithIndexFunction.apply(getFloat(i), i));
+ this.setFloat(i, mapWithIndexFunction.apply(this.getFloat(i), i));
}
return this;
}
diff --git a/source/net/yacy/ai/llama3/Tensor/Q4_0FloatTensor.java b/source/net/yacy/ai/llama3/Tensor/Q4_0FloatTensor.java
index a52a6df88..e413f2436 100644
--- a/source/net/yacy/ai/llama3/Tensor/Q4_0FloatTensor.java
+++ b/source/net/yacy/ai/llama3/Tensor/Q4_0FloatTensor.java
@@ -189,6 +189,7 @@ public final class Q4_0FloatTensor extends FloatTensor implements Tensor {
final int quantOffset = thisBlockOffset + QUANT_FLOAT16_BYTES;
float blockResult = 0.0f;
final int thatIndex = thatOffset + index;
+ /*
if (that instanceof ArrayFloatTensor) {
final ArrayFloatTensor thatArray = (ArrayFloatTensor) that;
final float[] b = thatArray.values;
@@ -198,12 +199,19 @@ public final class Q4_0FloatTensor extends FloatTensor implements Tensor {
final float valB1 = (float) FLOAT_ARRAY_HANDLE.get(b, thatIndex + i + QUANT_HALF_BLOCK);
blockResult += ((packed & 0x0F) - 8) * valB0 + (((packed >>> 4) & 0x0F) - 8) * valB1;
}
- } else {
+ } else if (that instanceof DirectBufferFloatTensor) {
+ final DirectBufferFloatTensor thatArray = (DirectBufferFloatTensor) that;
+ final FloatBuffer b = thatArray.floatBuffer;
+ for (int i = 0; i < QUANT_HALF_BLOCK; ++i) {
+ final byte packed = buffer.get(quantOffset + i);
+ blockResult += ((packed & 0x0F) - 8) * b.get(thatIndex + i) + (((packed >>> 4) & 0x0F) - 8) * b.get(thatIndex + i + QUANT_HALF_BLOCK);
+ }
+ } else {*/
for (int i = 0; i < QUANT_HALF_BLOCK; ++i) {
final byte packed = buffer.get(quantOffset + i);
blockResult += ((packed & 0x0F) - 8) * that.getFloat(thatIndex + i) + (((packed >>> 4) & 0x0F) - 8) * that.getFloat(thatIndex + i + QUANT_HALF_BLOCK);
}
- }
+ //}
result += blockResult * thisScale;
index += GGMLType.Q4_0.blockSize;
}
diff --git a/source/net/yacy/ai/llama3/Tensor/Tensor.java b/source/net/yacy/ai/llama3/Tensor/Tensor.java
index bfe75d6e7..dcbe32476 100644
--- a/source/net/yacy/ai/llama3/Tensor/Tensor.java
+++ b/source/net/yacy/ai/llama3/Tensor/Tensor.java
@@ -23,6 +23,7 @@
package net.yacy.ai.llama3.Tensor;
+import java.util.Arrays;
import java.util.function.IntConsumer;
import java.util.function.LongConsumer;
import java.util.stream.IntStream;
@@ -75,7 +76,12 @@ public interface Tensor {
public Tensor softmaxInPlace(final int thisOffset, final int size);
public Tensor saxpyInPlace(final int thisOffset, final Tensor that, final int thatOffset, final int size, final float a);
-
+
+ public static int numberOfElements(final int... dimensions) {
+ assert Arrays.stream(dimensions).allMatch(i -> i > 0);
+ return Arrays.stream(dimensions).reduce(Math::multiplyExact).orElseThrow();
+ }
+
public static void parallelFor(final int startInclusive, final int endExclusive, final IntConsumer action) {
if (startInclusive == 0 && endExclusive == 1) {
action.accept(0);