405 lines
11 KiB
C
405 lines
11 KiB
C
#include "cythoned/bridge.h"
|
|
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <mpi.h>
|
|
|
|
#define TAG_IDGAF 0
|
|
#define TAG_BATCH 1
|
|
#define TAG_NETWK 2
|
|
#define TAG_WEIGH 3
|
|
#define TAG_READY 4
|
|
#define TAG_BREAK 5
|
|
#define TAG_STLEN 6
|
|
#define TAG_SWORD 7
|
|
#define TAG_IWORD 8
|
|
#define TAG_INSTR 9
|
|
|
|
#define COMM 25
|
|
#define ITER 690
|
|
#define BS 32
|
|
#define EMB 32
|
|
#define WIN 2
|
|
#define FLPC 1
|
|
|
|
#define in_range(i, x) (size_t i = 0; i < (x); i++)
|
|
// I am honestly VERY sorry for this but power corrupts even the best of us
|
|
|
|
#define INFO_PRINTF(fmt, ...) \
|
|
do { fprintf(stderr, fmt, __VA_ARGS__); } while(0)
|
|
#define INFO_PRINTLN(what) \
|
|
do { fprintf(stderr, "%s\n", what); } while(0)
|
|
#define INFO_PRINT(what) \
|
|
do { fprintf(stderr, "%s", what); } while(0)
|
|
|
|
int g_argc; // sorry!
|
|
|
|
typedef enum{
|
|
TOKENIZER,
|
|
FILTERER,
|
|
BATCHER,
|
|
LEARNER,
|
|
VISUALIZER,
|
|
DISPATCHER
|
|
} Role;
|
|
|
|
int world_size() {
|
|
int n;
|
|
MPI_Comm_size(MPI_COMM_WORLD, &n);
|
|
return n;
|
|
}
|
|
|
|
int my_mpi_id() {
|
|
int i;
|
|
MPI_Comm_rank(MPI_COMM_WORLD, &i);
|
|
return i;
|
|
}
|
|
|
|
size_t number_of(Role what) {
|
|
switch (what) {
|
|
case TOKENIZER:
|
|
if (g_argc < 2) {
|
|
INFO_PRINTLN("NOT ENOUGH INPUTS!");
|
|
MPI_Abort(MPI_COMM_WORLD, 1);
|
|
}
|
|
return g_argc - 1;
|
|
case FILTERER:
|
|
return 1;
|
|
case BATCHER:
|
|
return 1;
|
|
case LEARNER:
|
|
return world_size()
|
|
- number_of(TOKENIZER)
|
|
- number_of(FILTERER)
|
|
- number_of(BATCHER)
|
|
- number_of(DISPATCHER)
|
|
- number_of(VISUALIZER);
|
|
case VISUALIZER:
|
|
return 0;
|
|
case DISPATCHER:
|
|
return 1;
|
|
}
|
|
}
|
|
|
|
int mpi_id_from_role_id(Role role, int rid) {
|
|
int base = 0;
|
|
for (Role r = TOKENIZER; r < role; r++) {
|
|
base += number_of(r);
|
|
}
|
|
return rid + base;
|
|
}
|
|
|
|
int role_id_from_mpi_id(Role role, int mid) {
|
|
int z = mpi_id_from_role_id(role, 0);
|
|
int rid = mid - z;
|
|
if (rid >= number_of(role) || rid < 0) {
|
|
INFO_PRINTF("%d is not a %d\n", mid, role);
|
|
MPI_Abort(MPI_COMM_WORLD, 1);
|
|
}
|
|
return rid;
|
|
}
|
|
|
|
Role map_node() {
|
|
int node = my_mpi_id();
|
|
size_t base = 0;
|
|
for (Role r = TOKENIZER; r <= DISPATCHER; r++) {
|
|
if (node < number_of(r) + base) return r;
|
|
base += number_of(r);
|
|
}
|
|
INFO_PRINTF("Something went wrong for node %d\n", node);
|
|
MPI_Abort(MPI_COMM_WORLD, 1); // this is bad
|
|
return -1; // Not going to happen anyway (i hope)
|
|
}
|
|
|
|
void announce_ready(int dest) {
|
|
int me = my_mpi_id();
|
|
MPI_Send(&me, 1, MPI_INT, dest, TAG_READY, MPI_COMM_WORLD);
|
|
}
|
|
|
|
int wait_for_ready() {
|
|
int ready;
|
|
MPI_Recv(&ready, 1, MPI_INT, MPI_ANY_SOURCE, TAG_READY, MPI_COMM_WORLD,
|
|
MPI_STATUS_IGNORE);
|
|
return ready;
|
|
}
|
|
|
|
void free_word(Word* w) {
|
|
free(w->data);
|
|
w->data = NULL;
|
|
w->mem = 0;
|
|
}
|
|
|
|
void free_wordlist(WordList* wl) {
|
|
for in_range(i, wl->mem) {
|
|
free_word(wl->words + i);
|
|
}
|
|
free(wl->words);
|
|
wl->words = NULL;
|
|
wl->n_words = 0;
|
|
}
|
|
|
|
void send_word(Word* w, int dest) {
|
|
long len = strlen(w->data);
|
|
MPI_Send(&len, 1, MPI_LONG, dest, TAG_STLEN, MPI_COMM_WORLD);
|
|
MPI_Send(w->data, len + 1, MPI_CHAR, dest, TAG_SWORD, MPI_COMM_WORLD);
|
|
}
|
|
|
|
void ssend_word(Word* w, int dest) {
|
|
long len = strlen(w->data);
|
|
MPI_Ssend(&len, 1, MPI_LONG, dest, TAG_STLEN, MPI_COMM_WORLD);
|
|
MPI_Ssend(w->data, len + 1, MPI_CHAR, dest, TAG_SWORD, MPI_COMM_WORLD);
|
|
}
|
|
|
|
int recv_word(Word* w, int src) {
|
|
long len;
|
|
MPI_Status stat;
|
|
MPI_Recv(&len, 1, MPI_LONG, src, TAG_STLEN, MPI_COMM_WORLD,
|
|
&stat);
|
|
int the_src = stat.MPI_SOURCE;
|
|
if (w->mem < len + 1) {
|
|
w->mem = len + 1;
|
|
w->data = realloc(w->data, sizeof(char) * w->mem);
|
|
}
|
|
MPI_Recv(w->data, len + 1, MPI_CHAR, the_src, TAG_SWORD, MPI_COMM_WORLD,
|
|
MPI_STATUS_IGNORE);
|
|
return the_src;
|
|
}
|
|
|
|
void tokenizer(const char* source) {
|
|
INFO_PRINTF("Starting tokenizer %d\n", getpid());
|
|
WordList wl = {0, 0, NULL};
|
|
size_t sync_ctr = 0;
|
|
while (get_tokens(&wl, source)) {
|
|
for in_range(i, wl.n_words) {
|
|
if (sync_ctr == 1000) {
|
|
ssend_word(&wl.words[i], mpi_id_from_role_id(FILTERER, 0));
|
|
sync_ctr = 0;
|
|
} else {
|
|
send_word(&wl.words[i], mpi_id_from_role_id(FILTERER, 0));
|
|
}
|
|
sync_ctr++;
|
|
}
|
|
}
|
|
}
|
|
|
|
void filterer() {
|
|
INFO_PRINTF("Starting filterer %d\n", getpid());
|
|
Word w = {0, NULL};
|
|
const size_t num_streams = number_of(TOKENIZER);
|
|
const size_t entry_size = 2 * WIN + 1;
|
|
const size_t bufsize = num_streams * entry_size;
|
|
|
|
long* buffer = malloc(bufsize * sizeof(long));
|
|
size_t* have = calloc(num_streams, sizeof(size_t));
|
|
|
|
int src = 0; // WLOG
|
|
while (1) {
|
|
int stream_offs;
|
|
while (have[src] != entry_size) {
|
|
src = recv_word(&w, MPI_ANY_SOURCE);
|
|
src = role_id_from_mpi_id(TOKENIZER, src);
|
|
stream_offs = src*entry_size;
|
|
buffer[stream_offs + have[src]] = vocab_idx_of(&w);
|
|
if (buffer[stream_offs + have[src]] != -1) have[src]++;
|
|
}
|
|
have[src] = 0;
|
|
MPI_Send(buffer + stream_offs, entry_size, MPI_LONG,
|
|
mpi_id_from_role_id(BATCHER, 0),
|
|
TAG_IWORD, MPI_COMM_WORLD);
|
|
}
|
|
free_word(&w);
|
|
free(buffer);
|
|
free(have);
|
|
}
|
|
|
|
void batcher() {
|
|
int s = 0;
|
|
const size_t entry_size = 2 * WIN + 1;
|
|
const size_t bufsize = BS * entry_size;
|
|
float* batch = malloc(bufsize * sizeof(float));
|
|
long* l_wid = malloc(entry_size * sizeof(long));
|
|
|
|
while (1) {
|
|
for in_range(r, BS) {
|
|
MPI_Recv(l_wid, entry_size, MPI_LONG,
|
|
mpi_id_from_role_id(FILTERER, 0),
|
|
TAG_IWORD, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
|
|
for in_range(c, entry_size) {
|
|
batch[r*entry_size + c] = (float)l_wid[c];
|
|
}
|
|
}
|
|
MPI_Recv(&s, 1, MPI_INT, MPI_ANY_SOURCE, TAG_READY, MPI_COMM_WORLD,
|
|
MPI_STATUS_IGNORE);
|
|
MPI_Send(batch, bufsize, MPI_FLOAT, s, TAG_BATCH, MPI_COMM_WORLD);
|
|
printf("!\n");
|
|
}
|
|
free(l_wid);
|
|
free(batch);
|
|
}
|
|
|
|
void free_weightlist(WeightList* wl) {
|
|
for in_range(i, wl->n_weights) {
|
|
free(wl->weights[i].shape);
|
|
free(wl->weights[i].W);
|
|
}
|
|
free(wl->weights);
|
|
}
|
|
|
|
void send_weights(const WeightList* wl, int dest) {
|
|
// This assumes that the receiving end knows exactly
|
|
// the number of elements being sent and has memory ready
|
|
// for it.
|
|
for in_range(i, wl->n_weights) {
|
|
long n_el = 1;
|
|
for in_range(k, wl->weights[i].dims) {
|
|
n_el *= wl->weights[i].shape[k];
|
|
}
|
|
MPI_Send(wl->weights[i].W, n_el, MPI_FLOAT, dest,
|
|
TAG_WEIGH, MPI_COMM_WORLD);
|
|
}
|
|
}
|
|
|
|
void recv_weights(WeightList* wl, int src) {
|
|
// This assumes that the sender sends stuff that is going
|
|
// to fit into memory in correct order too.
|
|
for in_range(i, wl->n_weights) {
|
|
long n_el = 1;
|
|
for in_range(d, wl->weights[i].dims) {
|
|
n_el *= wl->weights[i].shape[d];
|
|
}
|
|
MPI_Recv(wl->weights[i].W, n_el, MPI_FLOAT, src,
|
|
TAG_WEIGH, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
|
|
}
|
|
}
|
|
|
|
void learner() {
|
|
INFO_PRINTF("Starting learner %d\n", getpid());
|
|
int me = my_mpi_id();
|
|
|
|
PyObject* net = create_network(WIN, EMB);
|
|
create_test_dataset(WIN);
|
|
WeightList wl;
|
|
init_weightlist_like(&wl, net);
|
|
size_t entry_size = (2*WIN + 1);
|
|
size_t bufsize = BS * entry_size;
|
|
|
|
float* batch = malloc(bufsize * sizeof(float));
|
|
|
|
for in_range(i, COMM) {
|
|
recv_weights(&wl, mpi_id_from_role_id(DISPATCHER, 0));
|
|
set_net_weights(net, &wl);
|
|
for in_range(k, ITER) {
|
|
MPI_Send(&me, 1, MPI_INT, mpi_id_from_role_id(BATCHER, 0),
|
|
TAG_READY, MPI_COMM_WORLD);
|
|
MPI_Recv(batch, bufsize, MPI_FLOAT,
|
|
mpi_id_from_role_id(BATCHER, 0), TAG_BATCH, MPI_COMM_WORLD,
|
|
MPI_STATUS_IGNORE);
|
|
step_net(net, batch, BS);
|
|
}
|
|
update_weightlist(&wl, net);
|
|
send_weights(&wl, mpi_id_from_role_id(DISPATCHER, 0));
|
|
}
|
|
Py_DECREF(net);
|
|
free_weightlist(&wl);
|
|
free(batch);
|
|
}
|
|
|
|
void dispatcher() {
|
|
INFO_PRINTF("Starting dispatcher %d\n", getpid());
|
|
PyObject* frank = create_network(WIN, EMB);
|
|
create_test_dataset(WIN);
|
|
WeightList wl;
|
|
init_weightlist_like(&wl, frank);
|
|
update_weightlist(&wl, frank);
|
|
|
|
int lpr = number_of(LEARNER) * FLPC; // Learners per round
|
|
WeightList *wls = malloc(sizeof(WeightList) * lpr);
|
|
for in_range(i, lpr) {
|
|
init_weightlist_like(wls + i, frank);
|
|
}
|
|
int *round = malloc(sizeof(int) * lpr);
|
|
|
|
float first_loss = eval_net(frank);
|
|
float crt_loss = first_loss;
|
|
float min_loss = crt_loss;
|
|
time_t start = time(NULL);
|
|
for in_range(i, COMM) {
|
|
randidx(round, number_of(LEARNER), lpr);
|
|
|
|
for in_range(k, lpr) {
|
|
// INFO_PRINTF(" %5d", round[k]);
|
|
send_weights(&wl, mpi_id_from_role_id(LEARNER, round[k]));
|
|
}
|
|
// INFO_PRINTLN("");
|
|
for in_range(k, lpr) {
|
|
recv_weights(wls + k, mpi_id_from_role_id(LEARNER, round[k]));
|
|
}
|
|
combo_weights(&wl, wls, lpr);
|
|
set_net_weights(frank, &wl);
|
|
crt_loss = eval_net(frank);
|
|
min_loss = crt_loss < min_loss ? crt_loss : min_loss;
|
|
INFO_PRINTF("Round %ld, validation loss %f\n", i, crt_loss);
|
|
}
|
|
time_t finish = time(NULL);
|
|
float delta_t = finish - start;
|
|
float delta_l = first_loss - eval_net(frank);
|
|
INFO_PRINTF(
|
|
"Laptop MPI sgd consecutive_batch W%d E%d BS%d R%d bpe%d LPR%d,"
|
|
"%f,%f,%f\n", WIN, EMB, BS, COMM, ITER, lpr,
|
|
delta_l / COMM, delta_l / delta_t, min_loss);
|
|
Py_DECREF(frank);
|
|
free_weightlist(&wl);
|
|
for in_range(i, lpr) free_weightlist(wls + i);
|
|
free(wls);
|
|
free(round);
|
|
}
|
|
|
|
void visualizer() {
|
|
INFO_PRINTF("Starting visualizer %d\n", getpid());
|
|
serve();
|
|
}
|
|
|
|
int main (int argc, const char **argv) {
|
|
MPI_Init(NULL, NULL);
|
|
|
|
g_argc = argc;
|
|
|
|
// Cython Boilerplate
|
|
PyImport_AppendInittab("bridge", PyInit_bridge);
|
|
Py_Initialize();
|
|
PyRun_SimpleString("import sys\nsys.path.insert(0,'')");
|
|
PyObject* bridge_module = PyImport_ImportModule("bridge");
|
|
|
|
// Actual Code
|
|
int role_id;
|
|
switch (map_node()) {
|
|
case TOKENIZER:
|
|
role_id = role_id_from_mpi_id(TOKENIZER, my_mpi_id());
|
|
tokenizer(argv[role_id + 1]);
|
|
break;
|
|
case FILTERER:
|
|
filterer();
|
|
break;
|
|
case BATCHER:
|
|
batcher();
|
|
break;
|
|
case LEARNER:
|
|
learner();
|
|
break;
|
|
case DISPATCHER:
|
|
dispatcher();
|
|
break;
|
|
case VISUALIZER:
|
|
visualizer();
|
|
break;
|
|
default:
|
|
INFO_PRINTLN("DYING HORRIBLY!");
|
|
}
|
|
|
|
// Finalizing Boilerplate
|
|
Py_DECREF(bridge_module);
|
|
Py_Finalize();
|
|
MPI_Finalize();
|
|
}
|