301 lines
8.3 KiB
C
301 lines
8.3 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
#include "distributed.h"
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <stdarg.h>
|
|
|
|
/* ============================================================= */
|
|
/* Global State */
|
|
/* ============================================================= */
|
|
|
|
static fd_config g_config = {0};
|
|
static bool g_initialized = false;
|
|
|
|
/* ============================================================= */
|
|
/* Initialization */
|
|
/* ============================================================= */
|
|
|
|
fd_config fd_init(int world_size, int rank) {
|
|
g_config.world_size = world_size > 0 ? world_size : 1;
|
|
g_config.rank = rank >= 0 && rank < g_config.world_size ? rank : 0;
|
|
g_config.local_rank = g_config.rank; /* Simplified: assume single node */
|
|
g_config.is_main = (g_config.rank == 0);
|
|
g_initialized = true;
|
|
return g_config;
|
|
}
|
|
|
|
fd_config fd_init_from_env(void) {
|
|
int world_size = 1;
|
|
int rank = 0;
|
|
|
|
const char* ws = getenv("WORLD_SIZE");
|
|
const char* r = getenv("RANK");
|
|
|
|
if (ws) world_size = atoi(ws);
|
|
if (r) rank = atoi(r);
|
|
|
|
return fd_init(world_size, rank);
|
|
}
|
|
|
|
bool fd_is_initialized(void) {
|
|
return g_initialized;
|
|
}
|
|
|
|
fd_config fd_get_config(void) {
|
|
if (!g_initialized) {
|
|
return fd_init(1, 0);
|
|
}
|
|
return g_config;
|
|
}
|
|
|
|
void fd_cleanup(void) {
|
|
g_initialized = false;
|
|
memset(&g_config, 0, sizeof(g_config));
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Collective Operations */
|
|
/* ============================================================= */
|
|
|
|
void fd_all_reduce(tm_tensor* tensor) {
|
|
/* Single-process: no-op */
|
|
if (!tensor || g_config.world_size <= 1) return;
|
|
|
|
/* In multi-process mode, this would use MPI_Allreduce */
|
|
/* For now, just a placeholder */
|
|
}
|
|
|
|
void fd_all_reduce_mean(tm_tensor* tensor) {
|
|
if (!tensor) return;
|
|
|
|
fd_all_reduce(tensor);
|
|
|
|
/* Divide by world size */
|
|
if (g_config.world_size > 1) {
|
|
float* data = tm_data_f32(tensor);
|
|
float scale = 1.0f / g_config.world_size;
|
|
for (int64_t i = 0; i < tensor->numel; i++) {
|
|
data[i] *= scale;
|
|
}
|
|
}
|
|
}
|
|
|
|
void fd_broadcast(tm_tensor* tensor, int src_rank) {
|
|
/* Single-process: no-op */
|
|
if (!tensor || g_config.world_size <= 1) return;
|
|
(void)src_rank;
|
|
}
|
|
|
|
void fd_reduce(tm_tensor* tensor, int dst_rank) {
|
|
/* Single-process: no-op */
|
|
if (!tensor || g_config.world_size <= 1) return;
|
|
(void)dst_rank;
|
|
}
|
|
|
|
tm_tensor** fd_all_gather(tm_tensor* tensor) {
|
|
if (!tensor) return NULL;
|
|
|
|
tm_tensor** gathered = calloc(g_config.world_size, sizeof(tm_tensor*));
|
|
if (!gathered) return NULL;
|
|
|
|
/* Single-process: just clone the tensor */
|
|
gathered[0] = tm_clone(tensor);
|
|
|
|
/* In multi-process mode, would gather from all ranks */
|
|
for (int i = 1; i < g_config.world_size; i++) {
|
|
gathered[i] = tm_clone(tensor); /* Placeholder */
|
|
}
|
|
|
|
return gathered;
|
|
}
|
|
|
|
void fd_barrier(void) {
|
|
/* Single-process: no-op */
|
|
/* In multi-process mode: MPI_Barrier */
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Gradient Accumulation */
|
|
/* ============================================================= */
|
|
|
|
fd_grad_accum* fd_grad_accum_new(tm_tensor** params, int num_params,
|
|
int accumulation_steps) {
|
|
if (!params || num_params <= 0 || accumulation_steps <= 0) return NULL;
|
|
|
|
fd_grad_accum* accum = calloc(1, sizeof(fd_grad_accum));
|
|
if (!accum) return NULL;
|
|
|
|
accum->num_params = num_params;
|
|
accum->accumulation_steps = accumulation_steps;
|
|
accum->current_step = 0;
|
|
|
|
accum->grads = calloc(num_params, sizeof(tm_tensor*));
|
|
if (!accum->grads) {
|
|
free(accum);
|
|
return NULL;
|
|
}
|
|
|
|
/* Initialize gradient accumulators to zeros */
|
|
for (int i = 0; i < num_params; i++) {
|
|
if (params[i]) {
|
|
accum->grads[i] = tm_zeros(params[i]->shape, params[i]->ndim, params[i]->dtype);
|
|
}
|
|
}
|
|
|
|
return accum;
|
|
}
|
|
|
|
void fd_grad_accum_free(fd_grad_accum* accum) {
|
|
if (!accum) return;
|
|
|
|
if (accum->grads) {
|
|
for (int i = 0; i < accum->num_params; i++) {
|
|
tm_release(accum->grads[i]);
|
|
}
|
|
free(accum->grads);
|
|
}
|
|
|
|
free(accum);
|
|
}
|
|
|
|
void fd_grad_accum_step(fd_grad_accum* accum, tm_tensor** params) {
|
|
if (!accum || !params) return;
|
|
|
|
/* Accumulate gradients */
|
|
for (int i = 0; i < accum->num_params; i++) {
|
|
if (params[i] && params[i]->grad && accum->grads[i]) {
|
|
float* acc = tm_data_f32(accum->grads[i]);
|
|
float* grad = tm_data_f32(params[i]->grad);
|
|
|
|
for (int64_t j = 0; j < params[i]->grad->numel; j++) {
|
|
acc[j] += grad[j];
|
|
}
|
|
}
|
|
}
|
|
|
|
accum->current_step++;
|
|
}
|
|
|
|
bool fd_grad_accum_should_step(fd_grad_accum* accum) {
|
|
if (!accum) return false;
|
|
return accum->current_step >= accum->accumulation_steps;
|
|
}
|
|
|
|
void fd_grad_accum_finish(fd_grad_accum* accum, tm_tensor** params) {
|
|
if (!accum || !params) return;
|
|
|
|
float scale = 1.0f / accum->accumulation_steps;
|
|
|
|
/* Copy averaged gradients back to params */
|
|
for (int i = 0; i < accum->num_params; i++) {
|
|
if (params[i] && params[i]->grad && accum->grads[i]) {
|
|
float* dst = tm_data_f32(params[i]->grad);
|
|
float* src = tm_data_f32(accum->grads[i]);
|
|
|
|
for (int64_t j = 0; j < params[i]->grad->numel; j++) {
|
|
dst[j] = src[j] * scale;
|
|
}
|
|
}
|
|
}
|
|
|
|
fd_grad_accum_reset(accum);
|
|
}
|
|
|
|
void fd_grad_accum_reset(fd_grad_accum* accum) {
|
|
if (!accum) return;
|
|
|
|
accum->current_step = 0;
|
|
|
|
for (int i = 0; i < accum->num_params; i++) {
|
|
if (accum->grads[i]) {
|
|
memset(tm_data_f32(accum->grads[i]), 0,
|
|
accum->grads[i]->numel * sizeof(float));
|
|
}
|
|
}
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Data Parallel */
|
|
/* ============================================================= */
|
|
|
|
int* fd_shard_indices(int total_size, int* shard_size) {
|
|
if (total_size <= 0 || !shard_size) return NULL;
|
|
|
|
int base_size = total_size / g_config.world_size;
|
|
int remainder = total_size % g_config.world_size;
|
|
|
|
/* This rank gets base_size + 1 if rank < remainder */
|
|
int my_size = base_size + (g_config.rank < remainder ? 1 : 0);
|
|
*shard_size = my_size;
|
|
|
|
int* indices = malloc(my_size * sizeof(int));
|
|
if (!indices) return NULL;
|
|
|
|
/* Calculate starting index for this rank */
|
|
int start = 0;
|
|
for (int r = 0; r < g_config.rank; r++) {
|
|
start += base_size + (r < remainder ? 1 : 0);
|
|
}
|
|
|
|
for (int i = 0; i < my_size; i++) {
|
|
indices[i] = start + i;
|
|
}
|
|
|
|
return indices;
|
|
}
|
|
|
|
tm_tensor* fd_shard_batch(tm_tensor* batch, int batch_dim) {
|
|
(void)batch_dim;
|
|
if (!batch || g_config.world_size <= 1) {
|
|
return tm_retain(batch);
|
|
}
|
|
|
|
/* In multi-process mode, would slice the batch */
|
|
/* For single process, return the whole batch */
|
|
return tm_retain(batch);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Checkpointing */
|
|
/* ============================================================= */
|
|
|
|
bool fd_should_save(void) {
|
|
return g_config.is_main;
|
|
}
|
|
|
|
void fd_checkpoint_barrier(void) {
|
|
fd_barrier();
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Utility */
|
|
/* ============================================================= */
|
|
|
|
void fd_print(const char* fmt, ...) {
|
|
if (!g_config.is_main) return;
|
|
|
|
va_list args;
|
|
va_start(args, fmt);
|
|
vprintf(fmt, args);
|
|
va_end(args);
|
|
}
|
|
|
|
uint64_t fd_get_seed(uint64_t base_seed) {
|
|
return base_seed + g_config.rank * 1000;
|
|
}
|
|
|
|
void fd_print_config(fd_config* config) {
|
|
if (!config) {
|
|
printf("(null config)\n");
|
|
return;
|
|
}
|
|
|
|
printf("Distributed config:\n");
|
|
printf(" World size: %d\n", config->world_size);
|
|
printf(" Rank: %d\n", config->rank);
|
|
printf(" Local rank: %d\n", config->local_rank);
|
|
printf(" Is main: %s\n", config->is_main ? "true" : "false");
|
|
}
|