152 lines
4.6 KiB
C
152 lines
4.6 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
#ifndef FOUNDRY_DISTRIBUTED_H
|
|
#define FOUNDRY_DISTRIBUTED_H
|
|
|
|
#include <stdint.h>
|
|
#include <stdbool.h>
|
|
#include <stddef.h>
|
|
#include "../tensor/tensor.h"
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
/*
|
|
* Distributed training primitives.
|
|
* Replaces DeepSpeed/FSDP basics.
|
|
*
|
|
* This is a simplified single-process implementation that can be
|
|
* extended with MPI/NCCL for multi-node training.
|
|
*/
|
|
|
|
/* ============================================================= */
|
|
/* Types */
|
|
/* ============================================================= */
|
|
|
|
/* Process group */
|
|
typedef struct fd_group fd_group;
|
|
|
|
/* Distributed config */
|
|
typedef struct {
|
|
int world_size; /* Total number of processes */
|
|
int rank; /* This process's rank (0 to world_size-1) */
|
|
int local_rank; /* Rank within node */
|
|
bool is_main; /* Is this the main process (rank 0) */
|
|
} fd_config;
|
|
|
|
/* Gradient accumulation state */
|
|
typedef struct {
|
|
tm_tensor** grads; /* Accumulated gradients */
|
|
int num_params;
|
|
int accumulation_steps;
|
|
int current_step;
|
|
} fd_grad_accum;
|
|
|
|
/* ============================================================= */
|
|
/* Initialization */
|
|
/* ============================================================= */
|
|
|
|
/* Initialize distributed environment
|
|
* For single-process: world_size=1, rank=0
|
|
*/
|
|
fd_config fd_init(int world_size, int rank);
|
|
|
|
/* Initialize from environment variables (WORLD_SIZE, RANK, LOCAL_RANK) */
|
|
fd_config fd_init_from_env(void);
|
|
|
|
/* Check if initialized */
|
|
bool fd_is_initialized(void);
|
|
|
|
/* Get current config */
|
|
fd_config fd_get_config(void);
|
|
|
|
/* Cleanup */
|
|
void fd_cleanup(void);
|
|
|
|
/* ============================================================= */
|
|
/* Collective Operations */
|
|
/* ============================================================= */
|
|
|
|
/* All-reduce: sum gradients across all processes
|
|
* In single-process mode, this is a no-op
|
|
*/
|
|
void fd_all_reduce(tm_tensor* tensor);
|
|
|
|
/* All-reduce with averaging */
|
|
void fd_all_reduce_mean(tm_tensor* tensor);
|
|
|
|
/* Broadcast from rank 0 to all other ranks */
|
|
void fd_broadcast(tm_tensor* tensor, int src_rank);
|
|
|
|
/* Reduce to single rank */
|
|
void fd_reduce(tm_tensor* tensor, int dst_rank);
|
|
|
|
/* All-gather: gather tensors from all ranks */
|
|
tm_tensor** fd_all_gather(tm_tensor* tensor);
|
|
|
|
/* Barrier: synchronize all processes */
|
|
void fd_barrier(void);
|
|
|
|
/* ============================================================= */
|
|
/* Gradient Accumulation */
|
|
/* ============================================================= */
|
|
|
|
/* Create gradient accumulator */
|
|
fd_grad_accum* fd_grad_accum_new(tm_tensor** params, int num_params,
|
|
int accumulation_steps);
|
|
|
|
/* Free accumulator */
|
|
void fd_grad_accum_free(fd_grad_accum* accum);
|
|
|
|
/* Accumulate gradients from current backward pass */
|
|
void fd_grad_accum_step(fd_grad_accum* accum, tm_tensor** params);
|
|
|
|
/* Check if we should update weights */
|
|
bool fd_grad_accum_should_step(fd_grad_accum* accum);
|
|
|
|
/* Get averaged gradients and reset accumulator */
|
|
void fd_grad_accum_finish(fd_grad_accum* accum, tm_tensor** params);
|
|
|
|
/* Reset accumulator */
|
|
void fd_grad_accum_reset(fd_grad_accum* accum);
|
|
|
|
/* ============================================================= */
|
|
/* Data Parallel */
|
|
/* ============================================================= */
|
|
|
|
/* Shard dataset indices across ranks */
|
|
int* fd_shard_indices(int total_size, int* shard_size);
|
|
|
|
/* Get this rank's portion of a batch */
|
|
tm_tensor* fd_shard_batch(tm_tensor* batch, int batch_dim);
|
|
|
|
/* ============================================================= */
|
|
/* Checkpointing */
|
|
/* ============================================================= */
|
|
|
|
/* Only rank 0 should save checkpoints */
|
|
bool fd_should_save(void);
|
|
|
|
/* Synchronize before/after checkpointing */
|
|
void fd_checkpoint_barrier(void);
|
|
|
|
/* ============================================================= */
|
|
/* Utility */
|
|
/* ============================================================= */
|
|
|
|
/* Print on main process only */
|
|
void fd_print(const char* fmt, ...);
|
|
|
|
/* Get rank-specific seed */
|
|
uint64_t fd_get_seed(uint64_t base_seed);
|
|
|
|
/* Print distributed config */
|
|
void fd_print_config(fd_config* config);
|
|
|
|
#ifdef __cplusplus
|
|
}
|
|
#endif
|
|
|
|
#endif /* FOUNDRY_DISTRIBUTED_H */
|