mpi : trying to move more MPI stuff into ggml-mpi (WIP) (#2099)

This commit is contained in:
Georgi Gerganov 2023-07-09 14:08:53 +03:00
parent ef61acfbf5
commit 3232db628c
No known key found for this signature in database
GPG key ID: 449E073F9DC10735
11 changed files with 134 additions and 67 deletions

View file

@ -9,13 +9,31 @@ extern "C" {
struct ggml_tensor * ggml_mpi_send_tensor(
struct ggml_context * ctx,
struct ggml_tensor *src,
int dst_rank);
struct ggml_tensor * src,
int dst_rank);
struct ggml_tensor * ggml_mpi_recv_tensor(
struct ggml_context * ctx,
struct ggml_tensor *parent,
struct ggml_tensor *dst,
int src_rank);
struct ggml_tensor * parent,
struct ggml_tensor * dst,
int src_rank);
struct ggml_mpi_context;
void ggml_mpi_backend_init(void);
void ggml_mpi_backend_free(void);
struct ggml_mpi_context * ggml_mpi_init(void);
void ggml_mpi_free(struct ggml_mpi_context * ctx);
int ggml_mpi_rank(struct ggml_mpi_context * ctx);
struct ggml_tensor * ggml_mpi_eval_init(
struct ggml_mpi_context * ctx_mpi,
struct ggml_context * ctx,
int n_embd,
int * n_tokens,
int * n_past,
int * n_threads);
#ifdef __cplusplus
}