wehub-resource-sync
20 lines
542 B
C++
20 lines
542 B
C++
// Copyright (c) Microsoft Corporation.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
// DeepSpeed Team
|
|
|
|
#include "deepcompile.h"
|
|
|
|
#pragma once
|
|
|
|
namespace dc {
|
|
|
|
void register_graph_z1(long graph_id, const std::vector<long>& ds_ids);
|
|
void register_param(long ds_id,
|
|
const std::vector<int64_t>& ds_shape,
|
|
at::Tensor ds_tensor,
|
|
at::Tensor grad_buffer,
|
|
int64_t offset);
|
|
void update_param_grad_buffer(long ds_id, at::Tensor grad_buffer, int64_t offset);
|
|
} // namespace dc
|