Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement consistent tensor detach #7265

Merged
merged 3 commits into from
Jan 15, 2022
Merged

Implement consistent tensor detach #7265

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
bbbb83b
Implement consistent tensor detach
hjchen2 Jan 14, 2022
c3f1d01
Merge branch 'master' into dev_fix_consistent_nn_parameter
oneflow-ci-bot Jan 14, 2022
456b844
Merge branch 'master' into dev_fix_consistent_nn_parameter
oneflow-ci-bot Jan 15, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions oneflow/core/framework/tensor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,8 +99,8 @@ bool ConsistentTensor::is_cuda() const {
}

Maybe<Tensor> ConsistentTensor::detach() const {
std::shared_ptr<Tensor> t = std::make_shared<ConsistentTensor>(impl_);
return t;
std::shared_ptr<Tensor> tensor = std::make_shared<ConsistentTensor>(JUST(impl_->detach()));
return tensor;
}

Maybe<void> ConsistentTensor::set_data(const std::shared_ptr<Tensor>& other) {
Expand Down
13 changes: 13 additions & 0 deletions oneflow/core/framework/tensor_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,11 @@ size_t ConsistentTensorMeta::CalcHashValue() const {
^ std::hash<Symbol<ParallelDesc>>()(parallel_desc());
}

Maybe<ConsistentTensorImpl> LazyConsistentTensorImpl::detach() const {
auto detached_impl = std::make_shared<LazyConsistentTensorImpl>(tensor_meta_, false, true);
return std::shared_ptr<ConsistentTensorImpl>(detached_impl);
}

EagerConsistentTensorImpl::EagerConsistentTensorImpl(
Symbol<ConsistentTensorMeta> consistent_tensor_meta, bool requires_grad, bool is_leaf,
const std::shared_ptr<MirroredTensor>& cur_rank_phy_tensor)
Expand Down Expand Up @@ -272,5 +277,13 @@ Maybe<Shape> GetPhysicalShape(const Shape& logical_shape, const cfg::NdSbp& nd_s
return std::shared_ptr<EagerConsistentTensorImpl>(tensor_impl);
}

Maybe<ConsistentTensorImpl> EagerConsistentTensorImpl::detach() const {
auto detached_impl = JUST(EagerConsistentTensorImpl::New(tensor_meta_, false, true));
detached_impl->cur_rank_phy_tensor_ = cur_rank_phy_tensor_;
detached_impl->consumer_nd_sbp_constraint_ = consumer_nd_sbp_constraint_;
detached_impl->transport_token_ = transport_token_;
return std::shared_ptr<ConsistentTensorImpl>(detached_impl);
}

} // namespace one
} // namespace oneflow
6 changes: 6 additions & 0 deletions oneflow/core/framework/tensor_impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,8 @@ class ConsistentTensorImpl : public TensorImpl {
return Maybe<void>::Ok();
}

virtual Maybe<ConsistentTensorImpl> detach() const { RETURN_ERROR_WITH_BUG_PROMPT(); }

protected:
ConsistentTensorImpl(Symbol<ConsistentTensorMeta> tensor_meta, bool requires_grad, bool is_leaf)
: TensorImpl(requires_grad, is_leaf),
Expand Down Expand Up @@ -254,6 +256,8 @@ class LazyConsistentTensorImpl final : public ConsistentTensorImpl {

// Getters
bool is_lazy() const override { return true; }

Maybe<ConsistentTensorImpl> detach() const override;
};

class EagerConsistentTensorImpl final : public ConsistentTensorImpl {
Expand All @@ -277,6 +281,8 @@ class EagerConsistentTensorImpl final : public ConsistentTensorImpl {
const Optional<int64_t>& parallel_id,
bool requires_grad, bool is_leaf);

Maybe<ConsistentTensorImpl> detach() const override;

private:
EagerConsistentTensorImpl(Symbol<ConsistentTensorMeta> consistent_tensor_meta, bool requires_grad,
bool is_leaf,
Expand Down