MSCCL++ Integration (#22734)
Co-authored-by: Caio Rocha <caiorocha@microsof.com> Co-authored-by: empyreus <rjsouza1995@gmail.com>
This commit is contained in:
co-authored by
Caio Rocha
empyreus
parent
9c53031d2b
commit
c2eae96c56
@@ -89,20 +89,6 @@ std::tuple<std::vector<int64_t>, std::vector<int64_t>> get_graph_buffer_ipc_meta
|
||||
void register_buffer(fptr_t _fa, const std::vector<fptr_t>& fake_ipc_ptrs);
|
||||
void register_graph_buffers(
|
||||
fptr_t _fa, const std::vector<std::vector<int64_t>>& handles, const std::vector<std::vector<int64_t>>& offsets);
|
||||
|
||||
// mscclpp
|
||||
torch::Tensor mscclpp_generate_unique_id();
|
||||
fptr_t mscclpp_init_context(
|
||||
const torch::Tensor& unique_id,
|
||||
const int64_t rank,
|
||||
const int64_t world_size,
|
||||
torch::Tensor& scratch,
|
||||
torch::Tensor& put_buffer,
|
||||
const int64_t nranks_per_node,
|
||||
const std::vector<int64_t>& rank_to_node,
|
||||
const std::vector<int64_t>& rank_to_ib,
|
||||
const int64_t context_selection);
|
||||
void mscclpp_allreduce(fptr_t _context, torch::Tensor& inp, torch::Tensor& out, int64_t nthreads, int64_t nblocks);
|
||||
#endif
|
||||
|
||||
/*
|
||||
|
||||
Reference in New Issue
Block a user