llama_cpp 0.3.2 → 0.3.4
Sign up to get free protection for your applications and to get access to all the features.
- checksums.yaml +4 -4
- data/CHANGELOG.md +37 -0
- data/ext/llama_cpp/extconf.rb +9 -0
- data/ext/llama_cpp/llama_cpp.cpp +302 -112
- data/ext/llama_cpp/src/ggml-cuda.cu +677 -118
- data/ext/llama_cpp/src/ggml-metal.h +5 -1
- data/ext/llama_cpp/src/ggml-metal.m +65 -45
- data/ext/llama_cpp/src/ggml-metal.metal +610 -484
- data/ext/llama_cpp/src/ggml-mpi.c +216 -0
- data/ext/llama_cpp/src/ggml-mpi.h +39 -0
- data/ext/llama_cpp/src/ggml.c +1146 -812
- data/ext/llama_cpp/src/ggml.h +77 -19
- data/ext/llama_cpp/src/k_quants.h +8 -0
- data/ext/llama_cpp/src/llama.cpp +289 -104
- data/ext/llama_cpp/src/llama.h +46 -3
- data/lib/llama_cpp/version.rb +2 -2
- data/lib/llama_cpp.rb +2 -1
- data/sig/llama_cpp.rbs +14 -1
- metadata +4 -2
@@ -0,0 +1,216 @@
|
|
1
|
+
#include "ggml-mpi.h"
|
2
|
+
|
3
|
+
#include "ggml.h"
|
4
|
+
|
5
|
+
#include <mpi.h>
|
6
|
+
|
7
|
+
#include <stdio.h>
|
8
|
+
#include <stdlib.h>
|
9
|
+
|
10
|
+
#define MIN(a, b) ((a) < (b) ? (a) : (b))
|
11
|
+
|
12
|
+
#define UNUSED GGML_UNUSED
|
13
|
+
|
14
|
+
struct ggml_mpi_context {
|
15
|
+
int rank;
|
16
|
+
int size;
|
17
|
+
};
|
18
|
+
|
19
|
+
void ggml_mpi_backend_init(void) {
|
20
|
+
MPI_Init(NULL, NULL);
|
21
|
+
}
|
22
|
+
|
23
|
+
void ggml_mpi_backend_free(void) {
|
24
|
+
MPI_Finalize();
|
25
|
+
}
|
26
|
+
|
27
|
+
struct ggml_mpi_context * ggml_mpi_init(void) {
|
28
|
+
struct ggml_mpi_context * ctx = calloc(1, sizeof(struct ggml_mpi_context));
|
29
|
+
|
30
|
+
MPI_Comm_rank(MPI_COMM_WORLD, &ctx->rank);
|
31
|
+
MPI_Comm_size(MPI_COMM_WORLD, &ctx->size);
|
32
|
+
|
33
|
+
return ctx;
|
34
|
+
}
|
35
|
+
|
36
|
+
void ggml_mpi_free(struct ggml_mpi_context * ctx) {
|
37
|
+
free(ctx);
|
38
|
+
}
|
39
|
+
|
40
|
+
int ggml_mpi_rank(struct ggml_mpi_context * ctx) {
|
41
|
+
return ctx->rank;
|
42
|
+
}
|
43
|
+
|
44
|
+
void ggml_mpi_eval_init(
|
45
|
+
struct ggml_mpi_context * ctx_mpi,
|
46
|
+
int * n_tokens,
|
47
|
+
int * n_past,
|
48
|
+
int * n_threads) {
|
49
|
+
UNUSED(ctx_mpi);
|
50
|
+
|
51
|
+
// synchronize the worker node parameters with the root node
|
52
|
+
MPI_Barrier(MPI_COMM_WORLD);
|
53
|
+
|
54
|
+
MPI_Bcast(n_tokens, 1, MPI_INT, 0, MPI_COMM_WORLD);
|
55
|
+
MPI_Bcast(n_past, 1, MPI_INT, 0, MPI_COMM_WORLD);
|
56
|
+
MPI_Bcast(n_threads, 1, MPI_INT, 0, MPI_COMM_WORLD);
|
57
|
+
}
|
58
|
+
|
59
|
+
static int ggml_graph_get_node_idx(struct ggml_cgraph * gf, const char * name) {
|
60
|
+
struct ggml_tensor * t = ggml_graph_get_tensor(gf, name);
|
61
|
+
if (t == NULL) {
|
62
|
+
fprintf(stderr, "%s: tensor %s not found\n", __func__, name);
|
63
|
+
return -1;
|
64
|
+
}
|
65
|
+
|
66
|
+
for (int i = 0; i < gf->n_nodes; i++) {
|
67
|
+
if (gf->nodes[i] == t) {
|
68
|
+
return i;
|
69
|
+
}
|
70
|
+
}
|
71
|
+
|
72
|
+
fprintf(stderr, "%s: tensor %s not found in graph (should not happen)\n", __func__, name);
|
73
|
+
return -1;
|
74
|
+
}
|
75
|
+
|
76
|
+
static void ggml_mpi_tensor_send(struct ggml_tensor * t, int mpi_rank_dst) {
|
77
|
+
MPI_Datatype mpi_type;
|
78
|
+
|
79
|
+
switch (t->type) {
|
80
|
+
case GGML_TYPE_I32: mpi_type = MPI_INT32_T; break;
|
81
|
+
case GGML_TYPE_F32: mpi_type = MPI_FLOAT; break;
|
82
|
+
default: GGML_ASSERT(false && "not implemented");
|
83
|
+
}
|
84
|
+
|
85
|
+
const int retval = MPI_Send(t->data, ggml_nelements(t), mpi_type, mpi_rank_dst, 0, MPI_COMM_WORLD);
|
86
|
+
GGML_ASSERT(retval == MPI_SUCCESS);
|
87
|
+
}
|
88
|
+
|
89
|
+
static void ggml_mpi_tensor_recv(struct ggml_tensor * t, int mpi_rank_src) {
|
90
|
+
MPI_Datatype mpi_type;
|
91
|
+
|
92
|
+
switch (t->type) {
|
93
|
+
case GGML_TYPE_I32: mpi_type = MPI_INT32_T; break;
|
94
|
+
case GGML_TYPE_F32: mpi_type = MPI_FLOAT; break;
|
95
|
+
default: GGML_ASSERT(false && "not implemented");
|
96
|
+
}
|
97
|
+
|
98
|
+
MPI_Status status; UNUSED(status);
|
99
|
+
|
100
|
+
const int retval = MPI_Recv(t->data, ggml_nelements(t), mpi_type, mpi_rank_src, MPI_ANY_TAG, MPI_COMM_WORLD, &status);
|
101
|
+
GGML_ASSERT(retval == MPI_SUCCESS);
|
102
|
+
}
|
103
|
+
|
104
|
+
// TODO: there are many improvements that can be done to this implementation
|
105
|
+
void ggml_mpi_graph_compute_pre(
|
106
|
+
struct ggml_mpi_context * ctx_mpi,
|
107
|
+
struct ggml_cgraph * gf,
|
108
|
+
int n_layers) {
|
109
|
+
const int mpi_rank = ctx_mpi->rank;
|
110
|
+
const int mpi_size = ctx_mpi->size;
|
111
|
+
|
112
|
+
struct ggml_tensor * inp_tokens = ggml_graph_get_tensor(gf, "inp_tokens");
|
113
|
+
if (inp_tokens == NULL) {
|
114
|
+
fprintf(stderr, "%s: tensor 'inp_tokens' not found\n", __func__);
|
115
|
+
return;
|
116
|
+
}
|
117
|
+
|
118
|
+
struct ggml_tensor * inp0 = ggml_graph_get_tensor(gf, "layer_inp_0");
|
119
|
+
if (inp0 == NULL) {
|
120
|
+
fprintf(stderr, "%s: tensor 'inp0' not found\n", __func__);
|
121
|
+
return;
|
122
|
+
}
|
123
|
+
|
124
|
+
GGML_ASSERT(inp0 == gf->nodes[0]);
|
125
|
+
|
126
|
+
// distribute the compute graph into slices across the MPI nodes
|
127
|
+
//
|
128
|
+
// the main node (0) processes the last layers + the remainder of the compute graph
|
129
|
+
// and is responsible to pass the input tokens to the first node (1)
|
130
|
+
//
|
131
|
+
// node 1: [( 0) * n_per_node, ( 1) * n_per_node)
|
132
|
+
// node 2: [( 1) * n_per_node, ( 2) * n_per_node)
|
133
|
+
// ...
|
134
|
+
// node n-1: [(n-2) * n_per_node, (n-1) * n_per_node)
|
135
|
+
// node 0: [(n-1) * n_per_node, n_nodes)
|
136
|
+
//
|
137
|
+
if (mpi_rank > 0) {
|
138
|
+
if (mpi_rank == 1) {
|
139
|
+
// the first node (1) receives the input tokens from the main node (0)
|
140
|
+
ggml_mpi_tensor_recv(inp_tokens, 0);
|
141
|
+
} else {
|
142
|
+
// recv input data for each node into the "inp0" tensor (i.e. the first node in the compute graph)
|
143
|
+
ggml_mpi_tensor_recv(inp0, mpi_rank - 1);
|
144
|
+
}
|
145
|
+
} else if (mpi_size > 1) {
|
146
|
+
// node 0 sends the input tokens to node 1
|
147
|
+
ggml_mpi_tensor_send(inp_tokens, 1);
|
148
|
+
|
149
|
+
// recv the output data from the last node
|
150
|
+
ggml_mpi_tensor_recv(inp0, mpi_size - 1);
|
151
|
+
}
|
152
|
+
|
153
|
+
{
|
154
|
+
const int n_per_node = (n_layers + (mpi_size - 1)) / mpi_size;
|
155
|
+
|
156
|
+
const int mpi_idx = mpi_rank > 0 ? mpi_rank - 1 : mpi_size - 1;
|
157
|
+
|
158
|
+
const int il0 = (mpi_idx + 0) * n_per_node;
|
159
|
+
const int il1 = MIN(n_layers, (mpi_idx + 1) * n_per_node);
|
160
|
+
|
161
|
+
char name_l0[GGML_MAX_NAME];
|
162
|
+
char name_l1[GGML_MAX_NAME];
|
163
|
+
|
164
|
+
snprintf(name_l0, sizeof(name_l0), "layer_inp_%d", il0);
|
165
|
+
snprintf(name_l1, sizeof(name_l1), "layer_inp_%d", il1);
|
166
|
+
|
167
|
+
const int idx_l0 = ggml_graph_get_node_idx(gf, name_l0);
|
168
|
+
const int idx_l1 = mpi_rank > 0 ? ggml_graph_get_node_idx(gf, name_l1) + 1 : gf->n_nodes;
|
169
|
+
|
170
|
+
if (idx_l0 < 0 || idx_l1 < 0) {
|
171
|
+
fprintf(stderr, "%s: layer input nodes not found\n", __func__);
|
172
|
+
return;
|
173
|
+
}
|
174
|
+
|
175
|
+
// attach the input data to all nodes that need it
|
176
|
+
// TODO: not great - should be able to do this without modifying the compute graph (see next TODO below)
|
177
|
+
for (int i = idx_l0; i < idx_l1; i++) {
|
178
|
+
if (gf->nodes[i]->src[0] == gf->nodes[idx_l0]) {
|
179
|
+
gf->nodes[i]->src[0] = inp0;
|
180
|
+
}
|
181
|
+
if (gf->nodes[i]->src[1] == gf->nodes[idx_l0]) {
|
182
|
+
gf->nodes[i]->src[1] = inp0;
|
183
|
+
}
|
184
|
+
}
|
185
|
+
|
186
|
+
// TODO: instead of rearranging the nodes, we should be able to execute a subset of the compute graph
|
187
|
+
for (int i = 1; i < idx_l1 - idx_l0; i++) {
|
188
|
+
gf->nodes[i] = gf->nodes[idx_l0 + i];
|
189
|
+
gf->grads[i] = gf->grads[idx_l0 + i];
|
190
|
+
}
|
191
|
+
|
192
|
+
// the first node performs the "get_rows" operation, the rest of the nodes get the data from the previous node
|
193
|
+
if (mpi_idx != 0) {
|
194
|
+
gf->nodes[0]->op = GGML_OP_NONE;
|
195
|
+
}
|
196
|
+
|
197
|
+
gf->n_nodes = idx_l1 - idx_l0;
|
198
|
+
|
199
|
+
//fprintf(stderr, "%s: node %d: processing %d nodes [%d, %d)\n", __func__, mpi_rank, gf->n_nodes, il0, il1);
|
200
|
+
}
|
201
|
+
}
|
202
|
+
|
203
|
+
void ggml_mpi_graph_compute_post(
|
204
|
+
struct ggml_mpi_context * ctx_mpi,
|
205
|
+
struct ggml_cgraph * gf,
|
206
|
+
int n_layers) {
|
207
|
+
UNUSED(n_layers);
|
208
|
+
|
209
|
+
const int mpi_rank = ctx_mpi->rank;
|
210
|
+
const int mpi_size = ctx_mpi->size;
|
211
|
+
|
212
|
+
// send the output data to the next node
|
213
|
+
if (mpi_rank > 0) {
|
214
|
+
ggml_mpi_tensor_send(gf->nodes[gf->n_nodes - 1], (mpi_rank + 1) % mpi_size);
|
215
|
+
}
|
216
|
+
}
|
@@ -0,0 +1,39 @@
|
|
1
|
+
#pragma once
|
2
|
+
|
3
|
+
struct ggml_context;
|
4
|
+
struct ggml_tensor;
|
5
|
+
struct ggml_cgraph;
|
6
|
+
|
7
|
+
#ifdef __cplusplus
|
8
|
+
extern "C" {
|
9
|
+
#endif
|
10
|
+
|
11
|
+
struct ggml_mpi_context;
|
12
|
+
|
13
|
+
void ggml_mpi_backend_init(void);
|
14
|
+
void ggml_mpi_backend_free(void);
|
15
|
+
|
16
|
+
struct ggml_mpi_context * ggml_mpi_init(void);
|
17
|
+
void ggml_mpi_free(struct ggml_mpi_context * ctx);
|
18
|
+
|
19
|
+
int ggml_mpi_rank(struct ggml_mpi_context * ctx);
|
20
|
+
|
21
|
+
void ggml_mpi_eval_init(
|
22
|
+
struct ggml_mpi_context * ctx_mpi,
|
23
|
+
int * n_tokens,
|
24
|
+
int * n_past,
|
25
|
+
int * n_threads);
|
26
|
+
|
27
|
+
void ggml_mpi_graph_compute_pre(
|
28
|
+
struct ggml_mpi_context * ctx_mpi,
|
29
|
+
struct ggml_cgraph * gf,
|
30
|
+
int n_layers);
|
31
|
+
|
32
|
+
void ggml_mpi_graph_compute_post(
|
33
|
+
struct ggml_mpi_context * ctx_mpi,
|
34
|
+
struct ggml_cgraph * gf,
|
35
|
+
int n_layers);
|
36
|
+
|
37
|
+
#ifdef __cplusplus
|
38
|
+
}
|
39
|
+
#endif
|