diff options
| -rw-r--r-- | source/net/yacy/ai/llama3/Llama.java | 14 | ||||
| -rw-r--r-- | source/net/yacy/ai/llama3/Llama3.java | 1 | ||||
| -rw-r--r-- | source/net/yacy/ai/llama3/Model/GGMLTensorEntry.java | 10 | ||||
| -rw-r--r-- | source/net/yacy/ai/llama3/Model/GGUF.java | 8 | ||||
| -rw-r--r-- | source/net/yacy/ai/llama3/Model/ModelLoader.java | 8 | ||||
| -rw-r--r-- | source/net/yacy/ai/llama3/Tensor/ArrayFloatTensor.java | 2 | ||||
| -rw-r--r-- | source/net/yacy/ai/llama3/Tensor/DirectBufferFloatTensor.java | 68 | ||||
| -rw-r--r-- | source/net/yacy/ai/llama3/Tensor/FloatTensor.java | 26 | ||||
| -rw-r--r-- | source/net/yacy/ai/llama3/Tensor/Q4_0FloatTensor.java | 12 | ||||
| -rw-r--r-- | source/net/yacy/ai/llama3/Tensor/Tensor.java | 8 |
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); |
