diff --git a/csrc/compile/deepcompile.cpp b/csrc/compile/deepcompile.cpp index 5274ab887c34..1d49acdca511 100644 --- a/csrc/compile/deepcompile.cpp +++ b/csrc/compile/deepcompile.cpp @@ -4,6 +4,7 @@ // DeepSpeed Team #include "deepcompile.h" +#include "z3.h" #define USE_C10D_NCCL @@ -15,7 +16,8 @@ std::shared_ptr reduce_buckets = nullptr; c10::intrusive_ptr process_group = nullptr; c10::intrusive_ptr symm_mem = nullptr; -ncclComm_t nccl_comm; +ncclComm_t nccl_comm = nullptr; +bool nccl_comm_initialized = false; bool use_symm_mem; bool profile = false; bool pre_div_reduce = true; @@ -84,10 +86,21 @@ void reset() void cleanup() { reset(); + reset_z3_gather_buffer_pool(); + if (reduce_buckets) { + reduce_buckets->clear(); + reduce_buckets.reset(); + } + param_registry.reset(); - ncclCommDestroy(nccl_comm); + if (nccl_comm_initialized) { + ncclCommDestroy(nccl_comm); + nccl_comm = nullptr; + nccl_comm_initialized = false; + } process_group = nullptr; symm_mem = nullptr; + profile = false; } at::Tensor reduce_grad(at::Tensor grad_tensor, long graph_id, long ds_id) @@ -150,6 +163,7 @@ void init(c10::intrusive_ptr pg, // create a new nccl communicator std::memcpy(&ncclID, tensor.to(torch::Device(torch::kCPU)).data_ptr(), NCCL_UNIQUE_ID_BYTES); ncclCommInitRank(&nccl_comm, process_group->getSize(), ncclID, process_group->getRank()); + nccl_comm_initialized = true; param_registry = std::make_shared(); reduce_buckets = std::make_shared( diff --git a/csrc/compile/init.cpp b/csrc/compile/init.cpp index 7ed378fd3cdd..46930b1da9ee 100644 --- a/csrc/compile/init.cpp +++ b/csrc/compile/init.cpp @@ -105,6 +105,21 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) m.def("register_graph_z3", &dc::register_graph_z3, "Register graph with a list of ds parameter ids"); + m.def("set_z3_gather_buffer_pool_budget_for_test", + &dc::set_z3_gather_buffer_pool_budget_for_test, + "Override the adaptive ZeRO-3 gather-buffer-pool budget for tests"); + m.def("update_z3_gather_buffer_pool_allocator_pressure_for_test", + &dc::update_z3_gather_buffer_pool_allocator_pressure_for_test, + "Simulate allocator pressure for ZeRO-3 gather-buffer-pool tests"); + m.def("get_z3_gather_buffer_pool_state_for_test", + &dc::get_z3_gather_buffer_pool_state_for_test, + "Inspect ZeRO-3 gather-buffer-pool accounting for tests"); + m.def("set_z3_param_valid_for_test", + &dc::set_z3_param_valid_for_test, + "Set a registered ZeRO-3 parameter validity bit for tests"); + m.def("set_z3_prefetch_fail_after_exclusions_for_test", + &dc::set_z3_prefetch_fail_after_exclusions_for_test, + "Inject a prefetch preparation failure after admission exclusions for tests"); m.def("start_forward", &dc::start_forward, "Start forward pass"); m.def("end_forward", &dc::end_forward, "End forward pass"); m.def("start_backward", &dc::start_backward, "Start backward pass"); diff --git a/csrc/compile/z3.cpp b/csrc/compile/z3.cpp index 6c0c92c68d08..9a6cebbc646f 100644 --- a/csrc/compile/z3.cpp +++ b/csrc/compile/z3.cpp @@ -7,16 +7,638 @@ #include "deepcompile.h" #include +#include + +#include +#include +#include +#include +#include +#include namespace dc { const size_t TIMEOUT_SYMMETRIC_MEMORY_BARRIER = 60000; +namespace { + +class GatherBufferPool { +public: + enum class ReleaseDisposition { Retained, ResizeByCaller, Retired }; + + at::Tensor acquire(int64_t numel, + at::ScalarType dtype, + const c10::Device& device, + at::cuda::CUDAStream stream) + { + flushCompletedPressureRecovery(); + observeAllocatorPressure(device.index()); + if (!enabled_ || pressure_recovery_in_progress_) { return at::Tensor(); } + + Entry* best = nullptr; + for (auto& entry : entries_) { + if (entry.checked_out || entry.buffer.scalar_type() != dtype || + entry.buffer.device() != device || entry.buffer.numel() < numel) { + continue; + } + if (best == nullptr || entry.buffer.numel() < best->buffer.numel()) { best = &entry; } + } + if (best == nullptr) { return at::Tensor(); } + + best->ready_event->block(stream); + at::Tensor result = best->buffer.narrow(0, 0, numel); + best->checked_out = true; + best->ready_event.reset(); + best->last_use = ++clock_; + return result; + } + + void excludeFromAdmission(const at::Tensor& buffer) + { + if (!buffer.defined() || buffer.storage().nbytes() == 0) { return; } + non_retainable_storages_.insert(buffer.storage().unsafeGetStorageImpl()); + } + + void cancelAdmissionExclusion(const at::Tensor& buffer) + { + if (!buffer.defined()) { return; } + non_retainable_storages_.erase(buffer.storage().unsafeGetStorageImpl()); + } + + void discard(const at::Tensor& buffer) + { + if (!buffer.defined()) { return; } + + auto storage = buffer.storage(); + auto storage_impl = storage.unsafeGetStorageImpl(); + non_retainable_storages_.erase(storage_impl); + if (storage.nbytes() == 0) { + removeEntryOwnership(storage_impl, "discarded", false); + return; + } + + buffer.record_stream(at::cuda::getCurrentCUDAStream()); + removeEntryOwnership(storage_impl, "discarded", false); + } + + ReleaseDisposition release(const at::Tensor& buffer) + { + auto storage = buffer.storage(); + if (storage.nbytes() == 0) { + auto storage_impl = storage.unsafeGetStorageImpl(); + non_retainable_storages_.erase(storage_impl); + if (removeEntryOwnership(storage_impl, "released_zero_storage", false)) { + return ReleaseDisposition::Retired; + } + return ReleaseDisposition::ResizeByCaller; + } + + auto consumer_stream = at::cuda::getCurrentCUDAStream(); + buffer.record_stream(consumer_stream); + observeAllocatorPressure(buffer.device().index()); + + auto storage_impl = storage.unsafeGetStorageImpl(); + const bool excluded = non_retainable_storages_.erase(storage_impl) > 0; + + for (auto it = entries_.begin(); it != entries_.end(); ++it) { + if (storageImpl(*it) != storage_impl) { continue; } + if (excluded) { return retireEntry(it, storage_impl, "excluded", false); } + if (pressure_recovery_in_progress_) { + return retireEntry(it, storage_impl, "recovery_release", true); + } + if (!enabled_ || charged_bytes_ > budget_bytes_) { + return retireEntry(it, storage_impl, "evicted_on_return", true); + } + auto ready_event = std::make_shared(cudaEventDisableTiming); + ready_event->record(consumer_stream); + it->ready_event = ready_event; + it->checked_out = false; + it->last_use = ++clock_; + return ReleaseDisposition::Retained; + } + + if (excluded || !enabled_ || pressure_recovery_in_progress_) { + return ReleaseDisposition::ResizeByCaller; + } + + const size_t capacity_bytes = storage.nbytes(); + if (pressure_recovery_complete_ && + !hasRecoveryAdmissionSlot(capacity_bytes, buffer.scalar_type(), buffer.device())) { + return ReleaseDisposition::ResizeByCaller; + } + if (capacity_bytes > budget_bytes_ || !makeRoom(capacity_bytes)) { + return ReleaseDisposition::ResizeByCaller; + } + + const int64_t capacity_numel = static_cast(capacity_bytes / buffer.element_size()); + at::Tensor candidate = at::as_strided(buffer.detach(), {capacity_numel}, {1}, 0); + auto ready_event = std::make_shared(cudaEventDisableTiming); + ready_event->record(consumer_stream); + entries_.push_back({candidate, ready_event, false, ++clock_, capacity_bytes}); + charged_bytes_ += capacity_bytes; + if (charged_bytes_ > high_water_bytes_) { + high_water_bytes_ = charged_bytes_; + logState("high_water"); + } + return ReleaseDisposition::Retained; + } + + void setBudgetForTest(int64_t budget_bytes) + { + test_override_ = true; + enabled_ = budget_bytes > 0; + budget_bytes_ = budget_bytes > 0 ? static_cast(budget_bytes) : 0; + baseline_retries_.reset(); + idle_pressure_score_ = 0; + pressure_recovery_in_progress_ = false; + pressure_recovery_complete_ = false; + pressure_recovery_flush_pending_ = false; + pressure_recovery_budget_bytes_ = 0; + pressure_recovery_targets_.clear(); + makeRoom(0); + logState("test_budget"); + } + + void reset() + { + entries_.clear(); + non_retainable_storages_.clear(); + baseline_retries_.reset(); + idle_pressure_score_ = 0; + pressure_recovery_in_progress_ = false; + pressure_recovery_complete_ = false; + pressure_recovery_flush_pending_ = false; + pressure_recovery_budget_bytes_ = 0; + pressure_recovery_targets_.clear(); + budget_bytes_ = 0; + charged_bytes_ = 0; + high_water_bytes_ = 0; + clock_ = 0; + enabled_ = false; + test_override_ = false; + adaptive_budget_initialized_ = false; + } + + void observeAllocatorPressureForTest(int64_t retries, int64_t free_bytes, int64_t total_bytes) + { + TORCH_CHECK(!test_override_, + "allocator-pressure simulation requires the production adaptive pool"); + TORCH_CHECK(retries >= 0 && free_bytes >= 0 && total_bytes >= 0, + "allocator-pressure simulation values must be nonnegative"); + updateBudgetForAllocatorPressure( + retries, static_cast(free_bytes), static_cast(total_bytes)); + } + + std::vector stateForTest() const + { + size_t checked_out = 0; + for (const auto& entry : entries_) { checked_out += entry.checked_out ? 1 : 0; } + return {static_cast(budget_bytes_), + static_cast(charged_bytes_), + static_cast(high_water_bytes_), + static_cast(entries_.size()), + static_cast(checked_out), + baseline_retries_.value_or(-1), + enabled_ ? 1 : 0, + adaptive_budget_initialized_ ? 1 : 0, + idle_pressure_score_, + pressure_recovery_complete_ ? 1 : 0, + static_cast(pressure_recovery_budget_bytes_), + static_cast(recoveryPendingEntries()), + pressure_recovery_in_progress_ ? 1 : 0}; + } + +private: + struct Entry { + at::Tensor buffer; + std::shared_ptr ready_event; + bool checked_out; + uint64_t last_use; + size_t capacity_bytes; + }; + + struct RecoveryTarget { + size_t capacity_bytes; + at::ScalarType dtype; + c10::Device device; + uint64_t last_use; + }; + + static bool matchesRecoveryTarget(const Entry& entry, const RecoveryTarget& target) + { + return entry.capacity_bytes == target.capacity_bytes && + entry.buffer.scalar_type() == target.dtype && entry.buffer.device() == target.device; + } + + bool hasRecoveryAdmissionSlot(size_t capacity_bytes, + at::ScalarType dtype, + const c10::Device& device) const + { + const size_t target_count = + std::count_if(pressure_recovery_targets_.begin(), + pressure_recovery_targets_.end(), + [&](const RecoveryTarget& target) { + return capacity_bytes == target.capacity_bytes && + dtype == target.dtype && device == target.device; + }); + const size_t resident_count = + std::count_if(entries_.begin(), entries_.end(), [&](const Entry& entry) { + return entry.capacity_bytes == capacity_bytes && + entry.buffer.scalar_type() == dtype && entry.buffer.device() == device; + }); + return resident_count < target_count; + } + + size_t recoveryPendingEntries() const + { + std::vector matched(pressure_recovery_targets_.size(), false); + size_t resident_targets = 0; + for (const auto& entry : entries_) { + for (size_t i = 0; i < pressure_recovery_targets_.size(); ++i) { + if (!matched[i] && matchesRecoveryTarget(entry, pressure_recovery_targets_[i])) { + matched[i] = true; + ++resident_targets; + break; + } + } + } + return pressure_recovery_targets_.size() - resident_targets; + } + + static c10::StorageImpl* storageImpl(const Entry& entry) + { + return entry.buffer.storage().unsafeGetStorageImpl(); + } + + static bool recoveryTargetPriority(const RecoveryTarget& lhs, const RecoveryTarget& rhs) + { + if (lhs.capacity_bytes != rhs.capacity_bytes) { + return lhs.capacity_bytes > rhs.capacity_bytes; + } + return lhs.last_use > rhs.last_use; + } + + void boundRecoveryTargets(size_t hard_cap) + { + std::stable_sort(pressure_recovery_targets_.begin(), + pressure_recovery_targets_.end(), + recoveryTargetPriority); + std::vector bounded_targets; + bounded_targets.reserve(pressure_recovery_targets_.size()); + size_t bounded_bytes = 0; + for (const auto& target : pressure_recovery_targets_) { + if (target.capacity_bytes > hard_cap - bounded_bytes) { continue; } + bounded_targets.push_back(target); + bounded_bytes += target.capacity_bytes; + } + pressure_recovery_targets_ = std::move(bounded_targets); + pressure_recovery_budget_bytes_ = bounded_bytes; + } + + void dropRecoveryTarget(const Entry& entry) + { + auto target = std::find_if(pressure_recovery_targets_.begin(), + pressure_recovery_targets_.end(), + [&](const RecoveryTarget& candidate) { + return candidate.last_use == entry.last_use && + matchesRecoveryTarget(entry, candidate); + }); + if (target == pressure_recovery_targets_.end()) { return; } + pressure_recovery_budget_bytes_ -= target->capacity_bytes; + pressure_recovery_targets_.erase(target); + budget_bytes_ = std::min(budget_bytes_, pressure_recovery_budget_bytes_); + enabled_ = budget_bytes_ > 0; + } + + ReleaseDisposition retireEntry(std::vector::iterator entry, + c10::StorageImpl* storage_impl, + const char* event, + bool keep_recovery_target) + { + // The final consumer stream was recorded before this method. Resize + // must succeed before pool accounting forgets the live storage. + at::native::resize_bytes_cuda(storage_impl, 0); + if (pressure_recovery_in_progress_ && !keep_recovery_target) { dropRecoveryTarget(*entry); } + charged_bytes_ -= entry->capacity_bytes; + entries_.erase(entry); + if (pressure_recovery_in_progress_ && entries_.empty()) { + completePressureRecovery(); + } else { + logState(event); + } + return ReleaseDisposition::Retired; + } + + bool removeEntryOwnership(c10::StorageImpl* storage_impl, + const char* event, + bool keep_recovery_target) + { + for (auto it = entries_.begin(); it != entries_.end(); ++it) { + if (storageImpl(*it) != storage_impl) { continue; } + if (pressure_recovery_in_progress_ && !keep_recovery_target) { + dropRecoveryTarget(*it); + } + charged_bytes_ -= it->capacity_bytes; + entries_.erase(it); + if (pressure_recovery_in_progress_ && entries_.empty()) { + completePressureRecovery(); + } else { + logState(event); + } + return true; + } + return false; + } + + void observeAllocatorPressure(c10::DeviceIndex device) + { + if (test_override_) { return; } + + const auto stats = c10::cuda::CUDACachingAllocator::getDeviceStats(device); + size_t free_bytes = 0; + size_t total_bytes = 0; + if (baseline_retries_ && stats.num_alloc_retries > baseline_retries_.value()) { + C10_CUDA_CHECK(cudaMemGetInfo(&free_bytes, &total_bytes)); + } + updateBudgetForAllocatorPressure(stats.num_alloc_retries, free_bytes, total_bytes); + } + + void updateBudgetForAllocatorPressure(int64_t retries, size_t free_bytes, size_t total_bytes) + { + if (!baseline_retries_) { + baseline_retries_ = retries; + idle_pressure_score_ = 0; + return; + } + if (retries < baseline_retries_.value()) { + baseline_retries_ = retries; + idle_pressure_score_ = 0; + return; + } + if (retries == baseline_retries_.value()) { return; } + const int64_t retry_delta = retries - baseline_retries_.value(); + baseline_retries_ = retries; + + constexpr size_t granularity = 2 * 1024 * 1024; + // Bound prefetch-expanded overlap while preserving headroom for large + // non-gather allocations in the compiled graph. Allocator retries are + // also the signal that idle cached gathers may need to yield memory. + // Keep the byte-exact working set on isolated retries. On the first + // sustained-pressure wave, stop new retention, release idle entries, + // and retire checked-out entries as their final consumers return. Keep + // a byte-exact, typed recovery plan, preserving the complete hot set + // when it fits the hard cap. The lifecycle latch prevents later waves + // from draining it again and recreating per-step churn. + size_t hard_cap = total_bytes / 32; + hard_cap -= hard_cap % granularity; + size_t free_target = free_bytes / 4; + free_target -= free_target % granularity; + const size_t observed_budget = std::min(hard_cap, std::max(free_target, charged_bytes_)); + budget_bytes_ = adaptive_budget_initialized_ ? std::min(budget_bytes_, observed_budget) + : observed_budget; + if (pressure_recovery_in_progress_ || pressure_recovery_complete_) { + boundRecoveryTargets(hard_cap); + budget_bytes_ = std::max(budget_bytes_, pressure_recovery_budget_bytes_); + } + adaptive_budget_initialized_ = true; + enabled_ = budget_bytes_ > 0; + idle_pressure_score_ = + retry_delta > std::numeric_limits::max() - idle_pressure_score_ + ? std::numeric_limits::max() + : idle_pressure_score_ + retry_delta; + if (pressure_recovery_in_progress_) { + idle_pressure_score_ = 0; + } else if (!pressure_recovery_complete_ && + idle_pressure_score_ >= kIdlePressureEvictionThreshold && + charged_bytes_ >= budget_bytes_ && beginPressureRecovery(hard_cap)) { + return; + } else { + const size_t charged_before_hard_cap = charged_bytes_; + makeRoom(0); + if (charged_bytes_ < charged_before_hard_cap) { idle_pressure_score_ = 0; } + } + logState(enabled_ ? "pressure" : "disabled"); + } + + bool beginPressureRecovery(size_t hard_cap) + { + if (entries_.empty()) { return false; } + + std::vector recovery_targets; + recovery_targets.reserve(entries_.size()); + for (const auto& entry : entries_) { + recovery_targets.push_back(RecoveryTarget{entry.capacity_bytes, + entry.buffer.scalar_type(), + entry.buffer.device(), + entry.last_use}); + } + pressure_recovery_targets_ = std::move(recovery_targets); + boundRecoveryTargets(hard_cap); + + for (auto it = entries_.begin(); it != entries_.end();) { + if (it->checked_out) { + ++it; + continue; + } + // Relinquish the pool's active ownership before flushing free + // blocks. Later gathers can still reuse allocator-cached storage. + auto victim_storage = it->buffer.storage(); + at::native::resize_bytes_cuda(victim_storage.unsafeGetStorageImpl(), 0); + if (non_retainable_storages_.erase(victim_storage.unsafeGetStorageImpl()) > 0) { + dropRecoveryTarget(*it); + } + charged_bytes_ -= it->capacity_bytes; + it = entries_.erase(it); + } + + idle_pressure_score_ = 0; + budget_bytes_ = pressure_recovery_budget_bytes_; + enabled_ = budget_bytes_ > 0; + pressure_recovery_in_progress_ = !entries_.empty(); + pressure_recovery_complete_ = false; + if (entries_.empty()) { + completePressureRecovery(); + } else { + logState("recovering_pressure"); + } + return true; + } + + void completePressureRecovery() + { + pressure_recovery_in_progress_ = false; + pressure_recovery_complete_ = true; + pressure_recovery_flush_pending_ = true; + idle_pressure_score_ = 0; + budget_bytes_ = pressure_recovery_budget_bytes_; + enabled_ = budget_bytes_ > 0; + logState("recovered_pressure"); + } + + void flushCompletedPressureRecovery() + { + if (!pressure_recovery_flush_pending_) { return; } + c10::cuda::CUDACachingAllocator::emptyCache(); + pressure_recovery_flush_pending_ = false; + logState("flushed_recovered_pressure"); + } + + bool makeRoom(size_t capacity_bytes) + { + while (charged_bytes_ + capacity_bytes > budget_bytes_) { + auto victim = entries_.end(); + for (auto it = entries_.begin(); it != entries_.end(); ++it) { + if (it->checked_out) { continue; } + if (victim == entries_.end() || it->last_use < victim->last_use) { victim = it; } + } + if (victim == entries_.end()) { return false; } + + auto victim_storage = victim->buffer.storage(); + at::native::resize_bytes_cuda(victim_storage.unsafeGetStorageImpl(), 0); + charged_bytes_ -= victim->capacity_bytes; + entries_.erase(victim); + } + return true; + } + + void logState(const char* event) const + { + if (std::getenv("DEEPSPEED_ALLOCATOR_TELEMETRY") == nullptr) { return; } + size_t checked_out = 0; + for (const auto& entry : entries_) { checked_out += entry.checked_out ? 1 : 0; } + std::cout << "DEEPSPEED_Z3_GATHER_BUFFER_POOL event=" << event + << " budget_bytes=" << budget_bytes_ << " charged_bytes=" << charged_bytes_ + << " high_water_bytes=" << high_water_bytes_ << " entries=" << entries_.size() + << " checked_out=" << checked_out + << " idle_pressure_score=" << idle_pressure_score_ + << " pressure_recovery_in_progress=" << (pressure_recovery_in_progress_ ? 1 : 0) + << " pressure_recovery_complete=" << (pressure_recovery_complete_ ? 1 : 0) + << " pressure_recovery_budget_bytes=" << pressure_recovery_budget_bytes_ + << " pressure_recovery_pending_entries=" << recoveryPendingEntries() << std::endl; + } + + static constexpr int64_t kIdlePressureEvictionThreshold = 3; + std::vector entries_; + std::unordered_set non_retainable_storages_; + std::optional baseline_retries_; + int64_t idle_pressure_score_ = 0; + bool pressure_recovery_in_progress_ = false; + bool pressure_recovery_complete_ = false; + bool pressure_recovery_flush_pending_ = false; + size_t pressure_recovery_budget_bytes_ = 0; + std::vector pressure_recovery_targets_; + size_t budget_bytes_ = 0; + size_t charged_bytes_ = 0; + size_t high_water_bytes_ = 0; + uint64_t clock_ = 0; + bool enabled_ = false; + bool test_override_ = false; + bool adaptive_budget_initialized_ = false; +}; + +class AdmissionExclusionRollback { +public: + explicit AdmissionExclusionRollback(std::shared_ptr pool) + : pool_(std::move(pool)) + { + } + + AdmissionExclusionRollback(const AdmissionExclusionRollback&) = delete; + AdmissionExclusionRollback& operator=(const AdmissionExclusionRollback&) = delete; + + ~AdmissionExclusionRollback() noexcept + { + if (committed_) { return; } + for (const auto& buffer : buffers_) { + try { + pool_->cancelAdmissionExclusion(buffer); + } catch (...) { + // Destructors must not replace the original prefetch exception. + } + } + } + + void exclude(const at::Tensor& buffer) + { + // Retain the Tensor before installing its raw StorageImpl identity so + // allocation/event exceptions cannot leave an ABA-prone dangling key. + buffers_.push_back(buffer); + pool_->excludeFromAdmission(buffers_.back()); + } + + size_t size() const { return buffers_.size(); } + + void commit() + { + committed_ = true; + buffers_.clear(); + } + +private: + std::shared_ptr pool_; + std::vector buffers_; + bool committed_ = false; +}; + +std::weak_ptr weak_gather_buffer_pool; +std::optional gather_buffer_pool_test_budget; +int64_t prefetch_fail_after_exclusions_for_test = 0; + +std::shared_ptr get_gather_buffer_pool() +{ + auto pool = weak_gather_buffer_pool.lock(); + if (!pool) { + pool = std::make_shared(); + weak_gather_buffer_pool = pool; + if (gather_buffer_pool_test_budget) { + pool->setBudgetForTest(gather_buffer_pool_test_budget.value()); + } + } + return pool; +} + +} // namespace + +void set_z3_gather_buffer_pool_budget_for_test(int64_t budget_bytes) +{ + gather_buffer_pool_test_budget = budget_bytes; + if (auto pool = weak_gather_buffer_pool.lock()) { pool->setBudgetForTest(budget_bytes); } +} + +void update_z3_gather_buffer_pool_allocator_pressure_for_test(int64_t retries, + int64_t free_bytes, + int64_t total_bytes) +{ + get_gather_buffer_pool()->observeAllocatorPressureForTest(retries, free_bytes, total_bytes); +} + +std::vector get_z3_gather_buffer_pool_state_for_test() +{ + return get_gather_buffer_pool()->stateForTest(); +} + +void set_z3_param_valid_for_test(long ds_id, bool valid) { param_registry->setValid(ds_id, valid); } + +void set_z3_prefetch_fail_after_exclusions_for_test(int64_t count) +{ + TORCH_CHECK(count >= 0, "prefetch exclusion failure count must be nonnegative"); + prefetch_fail_after_exclusions_for_test = count; +} + +void reset_z3_gather_buffer_pool() +{ + if (auto pool = weak_gather_buffer_pool.lock()) { pool->reset(); } + weak_gather_buffer_pool.reset(); + gather_buffer_pool_test_budget.reset(); + prefetch_fail_after_exclusions_for_test = 0; +} + class Z3CustomOpExecutor : public CustomOpExecutor { public: Z3CustomOpExecutor(c10::intrusive_ptr process_group, std::shared_ptr param_registry, std::shared_ptr reduce_buckets, + std::shared_ptr gather_buffer_pool, std::vector ds_ids, ncclComm_t nccl_comm, at::cuda::CUDAStream ag_stream, @@ -33,6 +655,7 @@ class Z3CustomOpExecutor : public CustomOpExecutor { rs_stream, copy_stream, pre_div_reduce), + gather_buffer_pool_(gather_buffer_pool), ag_stream_(ag_stream), offload_stream_(offload_stream), reload_stream_(reload_stream) @@ -146,8 +769,12 @@ class Z3CustomOpExecutor : public CustomOpExecutor { if (existing.defined() && existing.numel() == padded_numel) { output_buf = existing; } } if (!output_buf.defined()) { - at::cuda::CUDAStreamGuard guard(ag_stream_); - output_buf = torch::empty({padded_numel}, ds_tensor.options().dtype(target_dtype)); + output_buf = gather_buffer_pool_->acquire( + padded_numel, target_dtype, ds_tensor.device(), ag_stream_); + if (!output_buf.defined()) { + at::cuda::CUDAStreamGuard guard(ag_stream_); + output_buf = torch::empty({padded_numel}, ds_tensor.options().dtype(target_dtype)); + } } assert(hasKey(ag_comp_done_events_, ds_id)); @@ -176,6 +803,7 @@ class Z3CustomOpExecutor : public CustomOpExecutor { } std::unordered_map output_bufs; + AdmissionExclusionRollback admission_exclusions(gather_buffer_pool_); for (const auto& [ds_id, dtype] : invalid_params) { const DSParam& param = param_registry_->getParam(ds_id); const at::Tensor& ds_tensor = param.getDSTensor(); @@ -187,12 +815,23 @@ class Z3CustomOpExecutor : public CustomOpExecutor { auto existing = param_registry_->getGatheredParam(ds_id); if (existing.defined() && existing.numel() == padded_numel) { output_bufs[ds_id] = existing; - continue; } } - auto target_dtype = dtype ? dtype.value() : ds_tensor.scalar_type(); - output_bufs[ds_id] = - torch::empty({padded_numel}, ds_tensor.options().dtype(target_dtype)); + if (!hasKey(output_bufs, ds_id)) { + auto target_dtype = dtype ? dtype.value() : ds_tensor.scalar_type(); + at::cuda::CUDAStreamGuard guard(ag_stream_); + output_bufs[ds_id] = + torch::empty({padded_numel}, ds_tensor.options().dtype(target_dtype)); + } + // Prefetch lifetimes are already controlled by the memory-aware scheduler. + // Bind the exclusion to the selected storage so a stale ds_id generation + // cannot admit a different demand-gather buffer into the independent pool. + admission_exclusions.exclude(output_bufs.at(ds_id)); + if (prefetch_fail_after_exclusions_for_test > 0 && + admission_exclusions.size() >= + static_cast(prefetch_fail_after_exclusions_for_test)) { + throw std::runtime_error("injected prefetch preparation failure"); + } } for (const auto& [ds_id, _] : invalid_params) { @@ -210,6 +849,7 @@ class Z3CustomOpExecutor : public CustomOpExecutor { for (const auto& [ds_id, _] : invalid_params) { ag_comm_done_events_[ds_id]->record(ag_stream_); } + admission_exclusions.commit(); } void releaseParam(long ds_id, long n_users) @@ -226,10 +866,14 @@ class Z3CustomOpExecutor : public CustomOpExecutor { if (gathered_param.defined()) { // gathered param is undefined while profiling auto storage = gathered_param.storage(); if (storage.nbytes() > 0) { - // Required so the caching allocator defers reuse for consumer-stream kernels - // queued behind wait_allgather. - gathered_param.record_stream(at::cuda::getCurrentCUDAStream()); - at::native::resize_bytes_cuda(storage.unsafeGetStorageImpl(), 0); + // Demand gathers may enter the byte-bounded pool. Prefetched gathers + // keep the scheduler's ownership boundary and use the original + // resize-to-zero path after their final consumer. + const auto release_disposition = gather_buffer_pool_->release(gathered_param); + if (release_disposition == + GatherBufferPool::ReleaseDisposition::ResizeByCaller) { + at::native::resize_bytes_cuda(storage.unsafeGetStorageImpl(), 0); + } } const auto options = gathered_param.options(); @@ -420,6 +1064,7 @@ class Z3CustomOpExecutor : public CustomOpExecutor { bool hasParam(long ds_id) const { return hasKey(has_acc_grad_, ds_id); } private: + std::shared_ptr gather_buffer_pool_; at::cuda::CUDAStream ag_stream_; at::cuda::CUDAStream offload_stream_; at::cuda::CUDAStream reload_stream_; @@ -475,6 +1120,7 @@ void register_graph_z3(long graph_id, const std::vector& ds_ids) executors[graph_id] = std::make_shared(process_group, param_registry, reduce_buckets, + get_gather_buffer_pool(), ds_ids, nccl_comm, get_ag_stream(), @@ -553,13 +1199,22 @@ void set_persistent(long ds_id) // Allocate buffer here // Memory fragmentation will be more severe if we allocate in forward/backward + auto gather_buffer_pool = get_gather_buffer_pool(); + bool gathered = false; for (auto& it : executors) { if (it.second->hasParam(ds_id)) { auto executor = getExecutor(it.first, executors); auto dtype = param_registry->getParam(ds_id).getDtype(); executor->allgatherParam(ds_id, dtype, symm_mem); + gathered = true; } } + // Selective unsharding owns persistent storage for the remainder of the + // compiled lifecycle. If the initial gather reused a demand-pool entry, + // transfer that storage out of pool accounting without freeing it so + // pressure recovery cannot wait forever for a release that persistent + // parameters intentionally never issue. + if (gathered) { gather_buffer_pool->discard(param_registry->getGatheredParam(ds_id)); } } void prefetch_params_fused(long graph_id, @@ -584,17 +1239,22 @@ void invalidate_gathered_param(long ds_id) const DSParam& param = param_registry->getParam(ds_id); if (param.isPersistent()) { return; } + auto gathered_param = param_registry->getGatheredParam(ds_id); + get_gather_buffer_pool()->discard(gathered_param); param_registry->unregisterGatheredParam(ds_id); param_registry->registerGatheredParam(ds_id, at::Tensor()); } void clear_all_gathered_params() { + auto gather_buffer_pool = get_gather_buffer_pool(); for (const auto& it : param_registry->getParams()) { long ds_id = it.first; const DSParam& param = param_registry->getParam(ds_id); if (param.isPersistent()) { continue; } if (param_registry->hasGatheredParam(ds_id)) { + auto gathered_param = param_registry->getGatheredParam(ds_id); + gather_buffer_pool->discard(gathered_param); param_registry->unregisterGatheredParam(ds_id); } } diff --git a/csrc/compile/z3.h b/csrc/compile/z3.h index d467321cde4f..7987a9149cd2 100644 --- a/csrc/compile/z3.h +++ b/csrc/compile/z3.h @@ -10,6 +10,14 @@ namespace dc { void register_graph_z3(long graph_id, const std::vector& ds_ids); +void set_z3_gather_buffer_pool_budget_for_test(int64_t budget_bytes); +void update_z3_gather_buffer_pool_allocator_pressure_for_test(int64_t retries, + int64_t free_bytes, + int64_t total_bytes); +std::vector get_z3_gather_buffer_pool_state_for_test(); +void set_z3_param_valid_for_test(long ds_id, bool valid); +void set_z3_prefetch_fail_after_exclusions_for_test(int64_t count); +void reset_z3_gather_buffer_pool(); void register_graph_ops_z3(long graph_id, const std::vector& op_names, const std::vector& n_args); diff --git a/deepspeed/compile/backend.py b/deepspeed/compile/backend.py index ea207f39f6a5..a13645fd32dd 100644 --- a/deepspeed/compile/backend.py +++ b/deepspeed/compile/backend.py @@ -23,15 +23,16 @@ import deepspeed.comm as dist from deepspeed.accelerator import get_accelerator +from deepspeed.utils.allocator_telemetry import record_empty_cache from .fx import add_free_activations from .graph_param import DSGraphParamManager from .profilers import ProfilingResult -from .profilers.graph_profile import MemoryProfilingInterpreter +from .profilers.graph_profile import MemoryProfilingInterpreter, is_profile_incomplete from .patch_compiled_func import patch_compiled_func, unpatch_compiled_func, get_backward_inputs from .util import get_input_nodes, get_activation_node_names, get_index_by_graph_id, get_deepcompile_handle, log_rank0, is_backend_inductor from .partitioner import get_wrapped_partitioner -from .inductor import register_custom_ops, patch_create_aot_dispatcher_function +from .inductor import register_custom_ops, patch_create_aot_dispatcher_function, deepcompile_z3_inductor_config_patch from .input_storage import InputStorage remaining_schedule = None @@ -77,6 +78,22 @@ def clear(self): fwd_real_inputs = [] +def cleanup_compiled_backward_state(frame_id=None, owned_frames=None): + """Release engine-owned process-global compiled-backward state.""" + if frame_id is None: + if owned_frames is None: + frames_needing_bwd.clear() + else: + frames_needing_bwd.difference_update(owned_frames) + owned_frames.clear() + else: + frames_needing_bwd.discard(frame_id) + if owned_frames is not None: + owned_frames.discard(frame_id) + if len(frames_needing_bwd) == 0: + unpatch_compiled_func() + + def register_compile_pass(name: str, opt_pass_fn, contract=None): from .passes.contract import register_pass_contract opt_passes[name] = opt_pass_fn @@ -97,7 +114,8 @@ def init_schedule(schedule): remaining_schedule = deque(schedule) -def launch_compile_passes(global_steps: int): +def launch_compile_passes(global_steps: int, owned_frames=None): + """Advance the pass schedule and discard state owned by the previous compile cycle.""" global next_pass_step, next_passes if len(remaining_schedule) > 0 and global_steps == remaining_schedule[0][0]: @@ -109,6 +127,7 @@ def launch_compile_passes(global_steps: int): graph_order_with_frame_id.clear() profiling_results.clear() param_manager.clear() + cleanup_compiled_backward_state(owned_frames=owned_frames) frames_partitioned.clear() @@ -135,12 +154,15 @@ def set_time_and_tensor_size(graph_id, graph: Graph, mem, bwd, profiling_results profiling_results[graph_id].fwd_mem_complete = mem_complete -def _sync_memory_profile_complete(profile_complete: bool) -> bool: +def _sync_memory_profile_complete(profile_complete: bool, process_group=None) -> bool: if not dist.is_initialized(): return profile_complete complete = torch.tensor([1 if profile_complete else 0], device=torch.device(get_accelerator().current_device())) - dist.all_reduce(complete, dist.ReduceOp.MIN) + if process_group is None: + dist.all_reduce(complete, dist.ReduceOp.MIN) + else: + dist.all_reduce(complete, dist.ReduceOp.MIN, group=process_group) return bool(complete.item()) @@ -151,16 +173,28 @@ def evaluate_symint_from_shape_env(sym_int_v): return sym_int_v.node.hint -def set_example_values_to_symints(real_inputs, param_indices=None): +_ZERO_PARAMETER_COMPILE_METADATA = ("ds_id", "ds_shape", "ds_persist", "ds_status", "ds_target_dtype") + + +def set_example_values_to_symints(real_inputs, param_indices=None, real_zero_params=None): real_inputs_ret = [] # Create a set of parameter indices for quick lookup param_idx_set = set() if param_indices is not None: param_idx_set = {i for i, _, _ in param_indices} + param_ds_ids = {i: ds_id for i, ds_id, _ in (param_indices or [])} + real_zero_params = real_zero_params or {} for i, v in enumerate(real_inputs): if isinstance(v, torch.Tensor): + real_zero_param = real_zero_params.get(param_ds_ids.get(i)) + if i in param_idx_set and real_zero_param is not None: + # Stored and fake profiling inputs both discard instance-bound + # ZeRO methods, so recover the original before materialization. + real_inputs_ret.append(real_zero_param) + continue + if is_fake(v): shape = [] for fs in v.shape: @@ -184,10 +218,12 @@ def set_example_values_to_symints(real_inputs, param_indices=None): # Create Parameter if this input index corresponds to a parameter if i in param_idx_set: - dummy_v = torch.nn.Parameter(dummy_v) - # Copy any additional attributes from the original if they exist - if hasattr(v, 'ds_id'): - dummy_v.ds_id = v.ds_id + dummy_v = torch.nn.Parameter(dummy_v, requires_grad=v.requires_grad) + # Profiling and graph-parameter consumers use these ZeRO + # attributes after symbolic fake inputs are materialized. + for attr in _ZERO_PARAMETER_COMPILE_METADATA: + if hasattr(v, attr): + setattr(dummy_v, attr, getattr(v, attr)) real_inputs_ret.append(dummy_v) else: @@ -201,6 +237,21 @@ def set_example_values_to_symints(real_inputs, param_indices=None): return tuple(real_inputs_ret) +def _get_fw_real_inputs(local_real_inputs, input_storage: InputStorage, graph_id: int, debug_log: bool = False): + """Resolve graph-local real inputs from the one-shot queue or persistent storage.""" + if local_real_inputs: + return local_real_inputs.popleft() + + if input_storage.has_data(): + if debug_log: + log_rank0(f"Retrieving real inputs from storage for graph_id={graph_id}", enable=True) + return input_storage.get() + + raise RuntimeError(f"No real inputs available for graph_id {graph_id}. " + f"Local queue size: {len(local_real_inputs)}, " + f"storage has data: {input_storage.has_data()}") + + def run_opt_passes(opt_passes: List[Callable], gm: GraphModule, graph_id: int, @@ -210,12 +261,14 @@ def run_opt_passes(opt_passes: List[Callable], mem_budget: float, param_manager, bwd: bool, - debug_log=False) -> None: + debug_log=False, + process_group=None) -> None: + """Apply scheduled graph passes and retain only complete post-pass memory profiles.""" with unset_fake_temporarily(): get_accelerator().synchronize() gc.collect() - get_accelerator().empty_cache() + record_empty_cache("backend.pre-opt-passes", get_accelerator().empty_cache) for i, opt_pass_fn in enumerate(opt_passes): log_rank0(f"Running opt pass {i} for graph {graph_id}. bwd={bwd}", enable=debug_log) @@ -227,23 +280,33 @@ def run_opt_passes(opt_passes: List[Callable], gm.graph.lint() gm.recompile() - mem_prof = MemoryProfilingInterpreter(gm, debug_log=debug_log) - mem_prof.run(*create_inputs_fn()) - profile_complete = _sync_memory_profile_complete(mem_prof.profile_complete) - if profile_complete: - mem = [(name, current_alloc, delta, peak) for name, current_alloc, delta, peak in mem_prof.mem_record] - else: + # Re-profiling an already incomplete graph would turn synthetic + # backfilled metadata into a seemingly valid memory profile. + operator_profile_complete = _sync_memory_profile_complete(not is_profile_incomplete(gm.graph), + process_group) + if not operator_profile_complete: + profile_complete = False mem = [] + else: + mem_prof = MemoryProfilingInterpreter(gm, debug_log=debug_log, process_group=process_group) + mem_prof.run(*create_inputs_fn()) + profile_complete = _sync_memory_profile_complete(mem_prof.profile_complete, process_group) + if profile_complete: + mem = [(name, current_alloc, delta, peak) + for name, current_alloc, delta, peak in mem_prof.mem_record] + else: + mem = [] + del mem_prof set_time_and_tensor_size(graph_id, gm.graph, mem, bwd, profiling_results, profile_complete) with unset_fake_temporarily(): get_accelerator().synchronize() gc.collect() - get_accelerator().empty_cache() + record_empty_cache("backend.post-profile", get_accelerator().empty_cache) -def make_backend(backend, compile_config, compile_kwargs={}): +def make_backend(backend, compile_config, compile_kwargs={}, process_group=None, owned_frames=None): register_custom_ops() @@ -251,6 +314,10 @@ def make_backend(backend, compile_config, compile_kwargs={}): debug_log = compile_config.debug_log free_activation = compile_config.free_activation and not is_backend_inductor(backend) + if owned_frames is None: + owned_frames = set() + owner_token = object() + def backend_fn(gm: GraphModule, real_inputs): graph_id = id(gm.graph) @@ -263,34 +330,39 @@ def backend_fn(gm: GraphModule, real_inputs): # This check cannot be placed here because autograd creates the fw/bw compiler callables before graph # partitioning. It is thus postponed to the point where the fw compiler is called. frame_id = gm.meta["dynamo_compile_id"].frame_id + frame_key = (owner_token, frame_id) graph_order_with_frame_id.add_graph(graph_id, frame_id) z3_partition = any(hasattr(v, "ds_id") for v in real_inputs) if z3_partition: param_indices = [(i, input_val.ds_id, input_val.ds_shape) for i, input_val in enumerate(real_inputs) if isinstance(input_val, torch.nn.Parameter)] + real_zero_params = { + input_val.ds_id: input_val + for input_val in real_inputs if isinstance(input_val, torch.nn.Parameter) + and hasattr(input_val, "all_gather") and hasattr(input_val, "partition") + } else: assert all(hasattr(v, "param_id") for v in real_inputs if isinstance(v, torch.nn.Parameter)), "All param inputs should have param_id" param_indices = [(i, input_val.param_id, input_val.shape) for i, input_val in enumerate(real_inputs) if isinstance(input_val, torch.nn.Parameter)] - - global fwd_real_inputs + real_zero_params = {} # Create an InputStorage instance for this specific graph # It will be captured by the make_fw_graph closure, eliminating the need for graph ID management input_storage = InputStorage(keep_int_input_tensors=compile_config.keep_int_input_tensors, keep_all_input_tensors=compile_config.keep_all_input_tensors) - # Store in both list (for backward compatibility) and storage (for persistence) + # Store in a closure-local queue and storage (for persistence). # The input_storage keeps tensor metadata to handle cases where # backend_fn is called once but make_fw_graph is called multiple times - fwd_real_inputs.append(real_inputs) + local_fwd_real_inputs = deque([real_inputs]) input_storage.put(real_inputs) global profiling_results if graph_id not in profiling_results: - profiling_results[graph_id] = ProfilingResult() + profiling_results[graph_id] = ProfilingResult(process_group=process_group) profiling_results[graph_id].param_indices = param_indices def make_fw_graph(gm, sample_inputs): @@ -304,21 +376,11 @@ def make_fw_graph(gm, sample_inputs): if needs_backward: if len(frames_needing_bwd) == 0: patch_compiled_func() - frames_needing_bwd.add(frame_id) - - # Try to get real_inputs from the list first, then from storage - if fwd_real_inputs: - real_inputs = fwd_real_inputs.pop(0) - elif input_storage.has_data(): - # Note: input_storage is captured from the enclosing backend_fn scope - # Materialize tensors from storage when list is empty - log_rank0(f"Retrieving real inputs from storage for graph_id={graph_id}", enable=debug_log) - real_inputs = input_storage.get() - else: - raise RuntimeError(f"No real inputs available for graph_id {graph_id}. " - f"List size: {len(fwd_real_inputs)}, Storage has data: {input_storage.has_data()}") + frames_needing_bwd.add(frame_key) + owned_frames.add(frame_key) - real_inputs = set_example_values_to_symints(real_inputs) + real_inputs = _get_fw_real_inputs(local_fwd_real_inputs, input_storage, graph_id, debug_log=debug_log) + real_inputs = set_example_values_to_symints(real_inputs, param_indices, real_zero_params=real_zero_params) param_manager[graph_id] = DSGraphParamManager(gm.graph, real_inputs, param_indices) @@ -333,7 +395,8 @@ def make_fw_graph(gm, sample_inputs): mem_budget=.0, # unused param_manager=param_manager, bwd=False, - debug_log=debug_log) + debug_log=debug_log, + process_group=process_group) opt_pass_times.append(("fwd", graph_index, graph_id, time.time() - time_start)) @@ -375,7 +438,8 @@ def make_bw_graph(gm, sample_inputs): mem_budget=.0, # unused param_manager=param_manager, bwd=True, - debug_log=debug_log) + debug_log=debug_log, + process_group=process_group) # assert graph_id in param_manager, f"Graph {graph_id} not found in param_manager" @@ -385,9 +449,7 @@ def make_bw_graph(gm, sample_inputs): add_free_activations(graph_id, gm.graph, get_activation_node_names(gm.graph, param_nodes_bw, non_param_input_names)) - frames_needing_bwd.remove(frame_id) - if len(frames_needing_bwd) == 0: - unpatch_compiled_func() + cleanup_compiled_backward_state(frame_key, owned_frames) log_rank0( f"Bwd end {graph_index} graph_id={graph_id} alloc_mem={get_accelerator().memory_allocated()} graph={gm.graph}", @@ -415,10 +477,16 @@ def compiler_fn(gm, sample_inputs): partition_fn=partition_fn) return torch._dynamo.optimize(**compile_kwargs)(aot_mod) elif backend == "inductor": - patch_create_aot_dispatcher_function(graph_id, z3_partition, make_fw_graph, make_bw_graph, real_inputs, - param_indices, param_manager, frame_id, frames_partitioned) - - return torch._inductor.compile(gm, real_inputs) + restore_aotautograd = patch_create_aot_dispatcher_function(graph_id, z3_partition, make_fw_graph, + make_bw_graph, real_inputs, param_indices, + param_manager, frame_id, frames_partitioned) + try: + with deepcompile_z3_inductor_config_patch(z3_partition): + return torch._inductor.compile(gm, real_inputs) + finally: + # AotAutograd.__init__ is process-global; never leak this + # graph-specific compiler wiring into a later compilation. + restore_aotautograd() raise ValueError(f"Unsupported backend {backend}") diff --git a/deepspeed/compile/inductor.py b/deepspeed/compile/inductor.py index 3c2fa02ba4f8..6b991aebdaa6 100644 --- a/deepspeed/compile/inductor.py +++ b/deepspeed/compile/inductor.py @@ -3,6 +3,7 @@ # DeepSpeed Team +from contextlib import nullcontext from typing import Set import torch @@ -21,6 +22,37 @@ from .graph_param import DSGraphParamManager from .partitioner import get_wrapped_partitioner +_DEEP_COMPILE_Z3_INDUCTOR_REDUCTION_CONFIG = { + "triton.mix_order_reduction": False, + "triton.persistent_reductions": False, +} + + +def deepcompile_z3_inductor_config_patch(enabled: bool): + """Disable reduction heuristics that create oversized kernels for DeepCompile ZeRO-3 graphs.""" + if not enabled: + return nullcontext() + + inductor = getattr(torch, "_inductor", None) + config = getattr(inductor, "config", None) + if config is None or not hasattr(config, "patch"): + return nullcontext() + + triton_config = getattr(config, "triton", None) + if triton_config is None: + return nullcontext() + + overrides = { + config_name: value + for config_name, value in _DEEP_COMPILE_Z3_INDUCTOR_REDUCTION_CONFIG.items() + if hasattr(triton_config, + config_name.split(".", 1)[1]) + } + if not overrides: + return nullcontext() + + return config.patch(overrides) + def _get_graphsafe_run_with_rng_state(): try: @@ -46,6 +78,7 @@ def _mark_output_never_reuse(out, *, enabled): def patch_compiler(original_compiler, dc_compiler, z3_partition: bool, graph_id, graph_param_manager, bwd: bool): + """Wrap an AOT compiler with DeepCompile rewrites and ZeRO-3 fake-shape repair.""" def wrapped_compiler(gm, fake_inputs): mod_graph = dc_compiler(gm, fake_inputs) @@ -80,7 +113,8 @@ def wrapped_compiler(gm, fake_inputs): else: patched_inputs = fake_inputs - return original_compiler(gm, patched_inputs) + with deepcompile_z3_inductor_config_patch(z3_partition): + return original_compiler(gm, patched_inputs) return wrapped_compiler @@ -144,36 +178,45 @@ def _patch_deepcompile_aot_kwargs(kwargs: dict, *, graph_id: int, z3_partition: def patch_create_aot_dispatcher_function(graph_id: int, z3_partition: bool, make_fw_graph, make_bw_graph, real_inputs, param_indices, param_manager, frame_id: int, frames_partitioned: Set[int]): + """Temporarily install graph-specific AOT compilers and return an idempotent restore callback.""" from torch._dynamo.backends.common import AotAutograd import functools - def patch_aotautograd(): - # Unpatch if it was already patched - if hasattr(AotAutograd, "__original_init"): - AotAutograd.__init__ = AotAutograd.__original_init - - original_init = AotAutograd.__init__ - - @functools.wraps(original_init) - def patched_init(self, **kwargs): - _patch_deepcompile_aot_kwargs(kwargs, - graph_id=graph_id, - z3_partition=z3_partition, - make_fw_graph=make_fw_graph, - make_bw_graph=make_bw_graph, - real_inputs=real_inputs, - param_indices=param_indices, - param_manager=param_manager, - frame_id=frame_id, - frames_partitioned=frames_partitioned) - - original_init(self, **kwargs) - - AotAutograd.__original_init = original_init - AotAutograd.__init__ = patched_init - - patch_aotautograd() + # The constructor patch is process-global. Replace the currently installed + # DeepCompile patch before taking ownership for this graph. + if hasattr(AotAutograd, "__original_init"): + AotAutograd.__init__ = AotAutograd.__original_init + delattr(AotAutograd, "__original_init") + + original_init = AotAutograd.__init__ + + @functools.wraps(original_init) + def patched_init(self, **kwargs): + _patch_deepcompile_aot_kwargs(kwargs, + graph_id=graph_id, + z3_partition=z3_partition, + make_fw_graph=make_fw_graph, + make_bw_graph=make_bw_graph, + real_inputs=real_inputs, + param_indices=param_indices, + param_manager=param_manager, + frame_id=frame_id, + frames_partitioned=frames_partitioned) + + original_init(self, **kwargs) + + AotAutograd.__original_init = original_init + AotAutograd.__init__ = patched_init + + def restore_aotautograd(): + """Restore only this invocation's patch without clobbering a newer owner.""" + if AotAutograd.__init__ is patched_init: + AotAutograd.__init__ = original_init + if getattr(AotAutograd, "__original_init", None) is original_init: + delattr(AotAutograd, "__original_init") + + return restore_aotautograd def register_custom_ops(): diff --git a/deepspeed/compile/init_z1.py b/deepspeed/compile/init_z1.py index 880494b94200..025e45d23f0d 100644 --- a/deepspeed/compile/init_z1.py +++ b/deepspeed/compile/init_z1.py @@ -4,6 +4,7 @@ # DeepSpeed Team import copy +from functools import partial import torch @@ -185,5 +186,9 @@ def release_grad_buffer(group_idx=None): init_schedule(schedule) - engine.launch_compile_passes = launch_compile_passes - return make_backend(backend, compile_config, compile_kwargs=compile_kwargs) + engine._deepcompile_owned_frames = set() + engine.launch_compile_passes = partial(launch_compile_passes, owned_frames=engine._deepcompile_owned_frames) + return make_backend(backend, + compile_config, + compile_kwargs=compile_kwargs, + owned_frames=engine._deepcompile_owned_frames) diff --git a/deepspeed/compile/init_z3.py b/deepspeed/compile/init_z3.py index 0e20c02ede2e..8dc5c8421e54 100644 --- a/deepspeed/compile/init_z3.py +++ b/deepspeed/compile/init_z3.py @@ -3,10 +3,14 @@ # DeepSpeed Team +from functools import partial +from threading import Lock + import torch from deepspeed import comm as dist from deepspeed.accelerator import get_accelerator +from deepspeed.utils.allocator_telemetry import record_empty_cache from deepspeed.runtime.zero.partition_parameters import InsertPostInitMethodToModuleSubClasses from deepspeed.runtime.zero.parameter_offload import DeepSpeedZeRoOffload @@ -19,6 +23,54 @@ WARMUP = 5 _MISSING = object() +_DYNAMO_CONFIG_NAMES = ("force_parameter_static_shapes", "force_nn_module_property_static_shapes") +_DYNAMO_CONFIG_OWNERS = {} +_DYNAMO_CONFIG_LOCK = Lock() + + +def _allow_dynamo_dynamic_parameter_shapes_for_z3(compile_kwargs): + """Acquire process-wide ZeRO-3 Dynamo config ownership and return its release callback.""" + dynamo = getattr(torch, "_dynamo", None) + if dynamo is None: + try: + import torch._dynamo as dynamo + except ImportError: + return None + + dynamo_config = getattr(dynamo, "config", None) + if dynamo_config is None: + return None + + owner_token = object() + config_key = id(dynamo_config) + with _DYNAMO_CONFIG_LOCK: + state = _DYNAMO_CONFIG_OWNERS.get(config_key) + if state is None or state["config"] is not dynamo_config: + previous_values = { + config_name: getattr(dynamo_config, config_name) + for config_name in _DYNAMO_CONFIG_NAMES if hasattr(dynamo_config, config_name) + } + if not previous_values: + return None + state = {"config": dynamo_config, "previous_values": previous_values, "owner_tokens": set()} + _DYNAMO_CONFIG_OWNERS[config_key] = state + state["owner_tokens"].add(owner_token) + for config_name in state["previous_values"]: + setattr(dynamo_config, config_name, False) + + def restore(): + with _DYNAMO_CONFIG_LOCK: + state = _DYNAMO_CONFIG_OWNERS.get(config_key) + if state is None or state["config"] is not dynamo_config or owner_token not in state["owner_tokens"]: + return + state["owner_tokens"].remove(owner_token) + if state["owner_tokens"]: + return + for config_name, previous_value in state["previous_values"].items(): + setattr(dynamo_config, config_name, previous_value) + del _DYNAMO_CONFIG_OWNERS[config_key] + + return restore def _resolve_expected_grad_dtype(param): @@ -40,7 +92,7 @@ def init_z3(engine, backend, compile_config, compile_kwargs, schedule=None): if use_opt and hasattr(optimizer, "ipg_buckets"): optimizer.ipg_buckets.clear() - get_accelerator().empty_cache() + record_empty_cache("init-z3.ipg-clear", get_accelerator().empty_cache) dc = get_deepcompile_handle() dc.init(engine.data_parallel_group, compile_config, engine.zero_reduce_bucket_size()) @@ -102,9 +154,22 @@ def set_grad_buffer(_is_gradient_accumulation_boundary): if move_opt_states in passes or move_opt_states_sync in passes: init_offload_opt_states(optimizer, dc) - engine.launch_compile_passes = launch_compile_passes + engine._deepcompile_owned_frames = set() + engine.launch_compile_passes = partial(launch_compile_passes, owned_frames=engine._deepcompile_owned_frames) patch_fake_tensor() torch._inductor.config.size_asserts = False - return make_backend(backend, compile_config, compile_kwargs=compile_kwargs) + previous_restore = getattr(engine, "_deepcompile_dynamo_config_restore", None) + if previous_restore is not None: + previous_restore() + del engine._deepcompile_dynamo_config_restore + restore_dynamo_config = _allow_dynamo_dynamic_parameter_shapes_for_z3(compile_kwargs) + if restore_dynamo_config is not None: + engine._deepcompile_dynamo_config_restore = restore_dynamo_config + + return make_backend(backend, + compile_config, + compile_kwargs=compile_kwargs, + process_group=engine.data_parallel_group, + owned_frames=engine._deepcompile_owned_frames) diff --git a/deepspeed/compile/list_schedule.py b/deepspeed/compile/list_schedule.py index bdc2a9c2fe49..56301c329c5d 100644 --- a/deepspeed/compile/list_schedule.py +++ b/deepspeed/compile/list_schedule.py @@ -4,9 +4,9 @@ # DeepSpeed Team from collections import defaultdict -from typing import List, Dict +from typing import List, Dict, Optional from copy import copy -from dataclasses import dataclass +from dataclasses import dataclass, replace import torch from torch.fx import Graph, Node @@ -258,6 +258,228 @@ def register_arg(n: Node): return ordered_nodes +SCHEDULER_MEMORY_MARGIN = 0.1 +SCHEDULER_BUDGET_DIAGNOSTICS_ATTR = "_deepcompile_scheduler_budget_diagnostics" + + +@dataclass(frozen=True) +class SchedulerMemoryBudget: + """Rank-consistent allowance for scheduler-managed gathered parameter buffers. + + ``max_gathered_bytes`` is the headroom remaining after any profiled + non-gathered peak, output reservation, and safety margin are removed, or + an explicit gather-residency cap when no complete memory profile exists. + """ + + max_gathered_bytes: int + source: str + available_mem: int = 0 + output_size: int = 0 + safety_margin: int = 0 + total_mem: int = 0 + profiled_non_gathered_peak_mem: int = 0 + + @classmethod + def minimum_gather_residency(cls, max_single_allgather_bytes: int): + """Limit an incomplete-profile schedule to one largest-gather-sized residency. + + This is not an estimate of available memory or of an unobserved + non-gather peak. It only constrains the memory the scheduler controls: + simultaneously resident gathered parameter buffers. + """ + max_single_allgather_bytes = int(max_single_allgather_bytes or 0) + if max_single_allgather_bytes <= 0: + return None + return cls(max_gathered_bytes=max_single_allgather_bytes, source="incomplete_profile_minimum_gather_residency") + + def clamped_to_minimum_gather_residency(self, max_single_allgather_bytes: int): + """Keep an enabled budget large enough for one unavoidable gather. + + A smaller budget cannot be satisfied by any schedule and forces the + scheduler into its over-budget fallback. This floor changes only the + gathered residency controlled by the scheduler; it does not revise the + profiled non-gather memory estimate. + """ + max_single_allgather_bytes = int(max_single_allgather_bytes or 0) + if max_single_allgather_bytes <= self.max_gathered_bytes: + return self + return replace(self, + max_gathered_bytes=max_single_allgather_bytes, + source=f"{self.source}_clamped_to_minimum_gather_residency") + + @classmethod + def from_profiled_non_gathered_peak(cls, total_mem: int, profiled_non_gathered_peak_mem: int, output_size: int): + """Reserve profiled non-gather memory before budgeting transient gathers.""" + if total_mem is None or total_mem <= 0 or profiled_non_gathered_peak_mem is None or profiled_non_gathered_peak_mem <= 0: + return None + total_mem = int(total_mem) + profiled_non_gathered_peak_mem = int(profiled_non_gathered_peak_mem) + output_size = max(0, int(output_size or 0)) + safety_margin = int(total_mem * SCHEDULER_MEMORY_MARGIN) + max_gathered_bytes = max(0, total_mem - profiled_non_gathered_peak_mem - output_size - safety_margin) + return cls(max_gathered_bytes=max_gathered_bytes, + source="profiled_non_gathered_peak_memory", + available_mem=max_gathered_bytes, + output_size=output_size, + safety_margin=safety_margin, + total_mem=total_mem, + profiled_non_gathered_peak_mem=profiled_non_gathered_peak_mem) + + +class _GatheredParamTracker: + """Simulate gathered-buffer residency for a committed schedule prefix. + + A gathered buffer remains live until the final release for its ``ds_id``; + repeated gathers while it is live only increase the tracked allocation. + """ + + def __init__(self, + release_expected: Dict[int, int], + live_bytes_by_ds_id: Optional[Dict[int, int]] = None, + release_seen_by_ds_id: Optional[Dict[int, int]] = None, + live_bytes: int = 0, + peak_bytes: int = 0): + self.release_expected = release_expected + self.live_bytes_by_ds_id = dict(live_bytes_by_ds_id or {}) + self.release_seen_by_ds_id = dict(release_seen_by_ds_id or {}) + self.live_bytes = live_bytes + self.peak_bytes = peak_bytes + + def copy(self, *, reset_peak: bool = False): + """Copy residency state, optionally starting a candidate-local peak.""" + return _GatheredParamTracker( + self.release_expected, + self.live_bytes_by_ds_id, + self.release_seen_by_ds_id, + self.live_bytes, + self.live_bytes if reset_peak else self.peak_bytes, + ) + + def apply(self, node: Node): + """Advance residency through one node without treating reduce_grad as a release.""" + if node.target == torch.ops.dc.allgather_param.default: + ds_id = _get_ds_id(node) + size = _allgather_allocation_bytes(node) + current_size = self.live_bytes_by_ds_id.get(ds_id) + if current_size is None: + self.live_bytes_by_ds_id[ds_id] = size + self.live_bytes += size + elif size > current_size: + self.live_bytes_by_ds_id[ds_id] = size + self.live_bytes += size - current_size + self.release_seen_by_ds_id[ds_id] = 0 + elif is_release_node(node): + ds_id = _get_ds_id(node) + if ds_id in self.live_bytes_by_ds_id: + release_seen = self.release_seen_by_ds_id.get(ds_id, 0) + 1 + release_expected = max(1, self.release_expected.get(ds_id, _release_n_users(node))) + if release_seen >= release_expected: + self.live_bytes -= self.live_bytes_by_ds_id.pop(ds_id) + self.release_seen_by_ds_id.pop(ds_id, None) + else: + self.release_seen_by_ds_id[ds_id] = release_seen + + self.peak_bytes = max(self.peak_bytes, self.live_bytes) + + +def _get_ds_id(node: Node): + return node.args[2] + + +def _release_n_users(node: Node): + if len(node.args) > 3: + return int(node.args[3]) + return 1 + + +def _dtype_element_size(dtype): + if dtype is None: + return None + return torch.empty((), dtype=dtype).element_size() + + +def allgather_allocation_bytes(tensor_size: int, dtype, world_size: int): + """Return the padded allocation size used by a process-group all-gather.""" + element_size = _dtype_element_size(dtype) + tensor_size = int(tensor_size) + if tensor_size <= 0 or element_size is None or element_size <= 0 or world_size <= 1: + return tensor_size + true_numel, remainder = divmod(tensor_size, element_size) + if remainder != 0: + return tensor_size + padded_per_rank = (true_numel + world_size - 1) // world_size + return padded_per_rank * world_size * element_size + + +def _allgather_allocation_bytes(node: Node): + return int(node.meta.get("allgather_allocation_bytes", node.meta.get("tensor_size", 0))) + + +def _allgather_ranking_bytes(node: Node): + """Use legacy logical bytes for deterministic candidate ranking.""" + return int(node.meta.get("tensor_size", 0)) + + +def max_possible_gathered_bytes(graph: Graph): + """Upper-bound simultaneous gather residency, counting each ``ds_id`` once.""" + gathered_bytes_by_ds_id = {} + for node in graph.nodes: + if node.target != torch.ops.dc.allgather_param.default: + continue + ds_id = _get_ds_id(node) + gathered_bytes_by_ds_id[ds_id] = max(gathered_bytes_by_ds_id.get(ds_id, 0), _allgather_allocation_bytes(node)) + return sum(gathered_bytes_by_ds_id.values()) + + +def _build_release_expected(nodes: List[Node]): + """Derive the final-release count for every gathered parameter interval.""" + release_expected = defaultdict(int) + release_counts = defaultdict(int) + for node in nodes: + if is_release_node(node): + ds_id = _get_ds_id(node) + release_expected[ds_id] = max(release_expected[ds_id], _release_n_users(node)) + release_counts[ds_id] += 1 + + for ds_id, release_count in release_counts.items(): + release_expected[ds_id] = max(release_expected[ds_id], release_count) + + return dict(release_expected) + + +def _simulate_path(tracker: _GatheredParamTracker, nodes: List[Node]): + peak_bytes, _ = _simulate_path_stats(tracker, nodes) + return peak_bytes + + +def _simulate_path_stats(tracker: _GatheredParamTracker, nodes: List[Node]): + """Return this candidate's peak and ending residency without mutating the prefix. + + The committed tracker keeps a cumulative diagnostic peak, but an earlier + overflow must not make every later candidate appear over budget after its + buffers have been released. + """ + candidate_tracker = tracker.copy(reset_peak=True) + for node in nodes: + candidate_tracker.apply(node) + return candidate_tracker.peak_bytes, candidate_tracker.live_bytes + + +def profiled_non_gathered_peak(graph: Graph, mem_records): + """Return a conservative peak for scheduler-budget headroom. + + Operator profiling invalidates gathered parameters between nodes, so a + source-order residency replay is not guaranteed to be present in each + recorded peak. Keep the observed peak intact rather than subtracting + hypothetical gathered residency and risking an unsafe budget. + """ + return max((int(peak_mem) for _, _, _, peak_mem in mem_records), default=0) + + +def _fits_budget(scheduler_budget: Optional[SchedulerMemoryBudget], peak_bytes: int): + return scheduler_budget is None or peak_bytes <= scheduler_budget.max_gathered_bytes + + @dataclass class AllgatherTask: node: Node @@ -270,6 +492,10 @@ class AllgatherTask: n_scheduled_ags: int schedule_until_ag: List[Node] schedule_until_free: List[Node] + schedule_until_ag_peak_mem: int + schedule_until_free_peak_mem: int + schedule_until_ag_live_mem: int + schedule_until_free_live_mem: int def _free_path_allgather_key(task: AllgatherTask): @@ -280,8 +506,102 @@ def _fallback_allgather_key(task: AllgatherTask): return (task.free_acc_mem, task.n_scheduled_ags, task.allgather_acc_mem, task.free_cost, task.node.name) -def fast_free_schedule(graph: Graph, available_mem: int, output_size: int, debug_log: bool) -> Graph: +def _select_allgather_task(runnable_ags: List[AllgatherTask], scheduler_budget: Optional[SchedulerMemoryBudget]): + """Select the best fitting release path, then the best fitting gather-only path.""" + ags_with_no_additional_ag = [ + ag for ag in runnable_ags + if ag.free_acc_mem == 0 and _fits_budget(scheduler_budget, ag.schedule_until_free_peak_mem) + ] + if len(ags_with_no_additional_ag) > 0: + next_ag = sorted(ags_with_no_additional_ag, key=_free_path_allgather_key)[0] + return next_ag, next_ag.schedule_until_free + + fallback_ags = [ag for ag in runnable_ags if _fits_budget(scheduler_budget, ag.schedule_until_ag_peak_mem)] + if len(fallback_ags) == 0: + return None, None + next_ag = sorted(fallback_ags, key=_fallback_allgather_key)[0] + return next_ag, next_ag.schedule_until_ag + + +def _rejected_budget_candidates(runnable_ags: List[AllgatherTask], scheduler_budget: Optional[SchedulerMemoryBudget]): + """Describe tasks for which neither scheduler-eligible path fits the budget.""" + if scheduler_budget is None: + return [] + rejected = [] + for task in runnable_ags: + path_options = list(_over_budget_path_options(task, scheduler_budget)) + if any(_fits_budget(scheduler_budget, peak_bytes) for _, _, peak_bytes, _ in path_options): + continue + _, _, peak_bytes, _ = min(path_options, + key=lambda option: + (max(0, option[2] - scheduler_budget.max_gathered_bytes), option[2])) + rejected.append({ + "node": task.node.name, + "schedule_until_ag_peak_mem": task.schedule_until_ag_peak_mem, + "schedule_until_free_peak_mem": task.schedule_until_free_peak_mem, + "over_budget_bytes": max(0, peak_bytes - scheduler_budget.max_gathered_bytes), + }) + return rejected + + +def _over_budget_path_options(task: AllgatherTask, scheduler_budget: SchedulerMemoryBudget): + if task.free_acc_mem == 0: + yield ("until_free", task.schedule_until_free, task.schedule_until_free_peak_mem, + task.schedule_until_free_live_mem) + return + yield ("until_ag", task.schedule_until_ag, task.schedule_until_ag_peak_mem, task.schedule_until_ag_live_mem) + yield ("until_free", task.schedule_until_free, task.schedule_until_free_peak_mem, + task.schedule_until_free_live_mem) + + +def _select_over_budget_allgather_task(runnable_ags: List[AllgatherTask], scheduler_budget: SchedulerMemoryBudget): + """Choose the smallest deterministic peak overage when no candidate fits.""" + options = [] + for task in runnable_ags: + for path, nodes, peak_bytes, live_bytes in _over_budget_path_options(task, scheduler_budget): + overage_bytes = max(0, peak_bytes - scheduler_budget.max_gathered_bytes) + options.append((overage_bytes, peak_bytes, live_bytes, task.free_acc_mem, task.n_scheduled_ags, + task.allgather_acc_mem, task.free_cost, task.node.name, path, task, nodes)) + *_, task, nodes = min(options) + return task, nodes + + +def _budget_overflow_diagnostic(scheduler_budget: SchedulerMemoryBudget, task: AllgatherTask, path: str, + live_gathered_bytes: int): + peak_bytes = task.schedule_until_free_peak_mem if path == "until_free" else task.schedule_until_ag_peak_mem + return { + "source": scheduler_budget.source, + "max_gathered_bytes": scheduler_budget.max_gathered_bytes, + "live_gathered_bytes": live_gathered_bytes, + "selected_candidate": task.node.name, + "path": path, + "candidate_allgather_bytes": task.allgathered_mem, + "candidate_peak_bytes": peak_bytes, + "over_budget_bytes": max(0, peak_bytes - scheduler_budget.max_gathered_bytes), + } + + +def fast_free_schedule(graph: Graph, + available_mem: int, + output_size: int, + debug_log: bool, + *, + scheduler_budget: Optional[SchedulerMemoryBudget] = None) -> Graph: + """Order a ZeRO-3 graph while limiting scheduler-managed gather residency. + + The optional budget is already reduced across ranks by the caller. With no + budget, this preserves the legacy candidate ordering. With a budget, only + candidate simulations that fit are considered until all dependency levels + are exhausted; a deterministic least-over-budget path guarantees progress. + """ node_to_last_use, user_to_last_uses = get_last_uses(graph) + diagnostics = { + "budget": scheduler_budget, + "budget_rejections": 0, + "budget_rejected_candidates": [], + "budget_overflows": [], + "selected": [], + } # check tensor size for node in graph.nodes: @@ -293,6 +613,8 @@ def fast_free_schedule(graph: Graph, available_mem: int, output_size: int, debug graph) unscheduled_ags = [n for n in unscheduled if n.target == torch.ops.dc.allgather_param.default] + gathered_tracker = (_GatheredParamTracker(_build_release_expected(list(graph.nodes))) + if scheduler_budget is not None else None) release_nodes = defaultdict(list) for n in unscheduled: @@ -303,6 +625,9 @@ def fast_free_schedule(graph: Graph, available_mem: int, output_size: int, debug for ag_node in unscheduled_ags: last_use = node_to_last_use[ag_node] required_nodes = get_node_requirements(last_use, scheduled) + # Dependency gather count is the primary scheduling tier. Searching + # higher tiers is allowed only when every candidate in a lower tier is + # infeasible under the shared memory budget. ag_nodes_in_path[ag_node] = set(n for n in required_nodes if n.target == torch.ops.dc.allgather_param.default) reduce_nodes = [n for n in unscheduled if n.target == torch.ops.dc.reduce_grad.default] @@ -325,9 +650,12 @@ def fast_free_schedule(graph: Graph, available_mem: int, output_size: int, debug ag_nodes_count = {ag_node: len(nodes) for ag_node, nodes in ag_nodes_in_path.items()} count_list = sorted(set(ag_nodes_count.values())) - runnable_ags = [] + over_budget_ags = [] + next_ag = None + nodes_to_schedule = None for ag_count in count_list: + runnable_ags = [] target_unscheduled_ags = [ag for ag in unscheduled_ags if ag_nodes_count[ag] == ag_count] for node in target_unscheduled_ags: @@ -343,11 +671,13 @@ def fast_free_schedule(graph: Graph, available_mem: int, output_size: int, debug allgather_cost = sum(n.meta["device_time"] for n in schedule_until_ag) free_cost = sum(n.meta["device_time"] for n in diff_required_nodes) - allgathered_mem = node.meta["tensor_size"] - allgather_acc_mem = sum(n.meta["tensor_size"] for n in schedule_until_ag - if n.target == torch.ops.dc.allgather_param.default) - free_acc_mem = sum(n.meta["tensor_size"] for n in diff_required_nodes - if n.target == torch.ops.dc.allgather_param.default) + allgathered_mem = _allgather_allocation_bytes(node) + allgather_acc_mem = sum( + _allgather_ranking_bytes(n) for n in schedule_until_ag + if n.target == torch.ops.dc.allgather_param.default) + free_acc_mem = sum( + _allgather_ranking_bytes(n) for n in diff_required_nodes + if n.target == torch.ops.dc.allgather_param.default) schedule_until_free = schedule_until_ag + diff_required_nodes for release_node in release_nodes[ds_id]: @@ -358,72 +688,124 @@ def fast_free_schedule(graph: Graph, available_mem: int, output_size: int, debug n_scheduled_ags = len( [n for n in schedule_until_free if n.target == torch.ops.dc.allgather_param.default]) + if scheduler_budget is not None: + # Candidate simulation starts from the residency committed + # by earlier iterations and never mutates that prefix. + schedule_until_ag_peak_mem, schedule_until_ag_live_mem = _simulate_path_stats( + gathered_tracker, schedule_until_ag) + schedule_until_free_peak_mem, schedule_until_free_live_mem = _simulate_path_stats( + gathered_tracker, schedule_until_free) + else: + schedule_until_ag_peak_mem = 0 + schedule_until_free_peak_mem = 0 + schedule_until_ag_live_mem = 0 + schedule_until_free_live_mem = 0 + task = AllgatherTask(node, allgather_cost, free_cost, allgathered_mem, allgather_acc_mem, free_acc_mem, - last_use, n_scheduled_ags, schedule_until_ag, schedule_until_free) + last_use, n_scheduled_ags, schedule_until_ag, schedule_until_free, + schedule_until_ag_peak_mem, schedule_until_free_peak_mem, + schedule_until_ag_live_mem, schedule_until_free_live_mem) # print(f" ag_count {ag_count} allgather runnable {i}: {node} last_use: {node_to_last_use[node]} t: {t2-t1:.2f}") runnable_ags.append(task) if len(runnable_ags) > 0: - break - - assert len(runnable_ags) > 0, "No runnable allgather nodes" - - # Criteria of the choice: - # We want to choose allgather that does not require additional allgather until releasing the param. - # When we can find such a node, free_acc_mem will be zero. In that case, we choose the one with the smallest cost until free to minimize the period of occupying memory for the gathered param. - # If there is no such node, we choose the one with the smallest free_cost to minimize the period of occupying memory for the gathered param. - ags_with_no_additional_ag = [ag for ag in runnable_ags if ag.free_acc_mem == 0] - if len(ags_with_no_additional_ag) > 0: - sorted_ags = sorted(ags_with_no_additional_ag, key=_free_path_allgather_key) - next_ag = sorted_ags[0] - assert not debug_log or next_ag.free_acc_mem == 0 - nodes_to_schedule = next_ag.schedule_until_free - else: - # sorted_ags = sorted(runnable_ags, key=lambda x: x.allgathered_mem) - sorted_ags = sorted(runnable_ags, key=_fallback_allgather_key) - next_ag = sorted_ags[0] - nodes_to_schedule = next_ag.schedule_until_ag + if scheduler_budget is None: + next_ag, nodes_to_schedule = _select_allgather_task(runnable_ags, None) + break + else: + rejected = _rejected_budget_candidates(runnable_ags, scheduler_budget) + diagnostics["budget_rejections"] += len(rejected) + diagnostics["budget_rejected_candidates"].extend(rejected) + next_ag, nodes_to_schedule = _select_allgather_task(runnable_ags, scheduler_budget) + if next_ag is not None: + assert not debug_log or nodes_to_schedule is not next_ag.schedule_until_free or next_ag.free_acc_mem == 0 + break + over_budget_ags.extend(runnable_ags) + + if next_ag is None: + if scheduler_budget is not None and len(over_budget_ags) > 0: + # Failing compilation cannot improve memory pressure. Commit + # the deterministic path with the smallest peak overage and + # retain the overflow in diagnostics instead. + next_ag, nodes_to_schedule = _select_over_budget_allgather_task(over_budget_ags, scheduler_budget) + else: + raise AssertionError("No runnable allgather nodes") # print(f" next_ag {next_ag}") for n in nodes_to_schedule: + # Only the selected path advances the committed residency model; + # rejected simulations operated on tracker copies above. scheduled.append(n) unscheduled.remove(n) - - unscheduled_ags.remove(next_ag.node) - - ag_nodes_in_path.pop(next_ag.node) - for ag_node, nodes in ag_nodes_in_path.items(): - if next_ag.node in nodes: - nodes.remove(next_ag.node) + if gathered_tracker is not None: + gathered_tracker.apply(n) + + scheduled_ags = [n for n in nodes_to_schedule if n.target == torch.ops.dc.allgather_param.default] + for scheduled_ag in scheduled_ags: + if scheduled_ag in unscheduled_ags: + unscheduled_ags.remove(scheduled_ag) + + ag_nodes_in_path.pop(scheduled_ag, None) + for ag_node, nodes in ag_nodes_in_path.items(): + if scheduled_ag in nodes: + nodes.remove(scheduled_ag) + + selected_diagnostic = { + "node": next_ag.node.name, + "path": "until_free" if nodes_to_schedule is next_ag.schedule_until_free else "until_ag", + } + if scheduler_budget is not None: + selected_diagnostic.update({ + "live_gathered_bytes": gathered_tracker.live_bytes, + "peak_gathered_bytes": gathered_tracker.peak_bytes, + "schedule_until_ag_peak_mem": next_ag.schedule_until_ag_peak_mem, + "schedule_until_free_peak_mem": next_ag.schedule_until_free_peak_mem, + "schedule_until_ag_live_mem": next_ag.schedule_until_ag_live_mem, + "schedule_until_free_live_mem": next_ag.schedule_until_free_live_mem, + }) + diagnostics["selected"].append(selected_diagnostic) + if scheduler_budget is not None: + selected_peak_bytes = (next_ag.schedule_until_free_peak_mem if nodes_to_schedule + is next_ag.schedule_until_free else next_ag.schedule_until_ag_peak_mem) + if selected_peak_bytes > scheduler_budget.max_gathered_bytes: + diagnostics["budget_overflows"].append( + _budget_overflow_diagnostic(scheduler_budget, next_ag, diagnostics["selected"][-1]["path"], + diagnostics["selected"][-1]["live_gathered_bytes"])) # Schedule reduce nodes when possible to free memory earlier reduces_to_schedule = [] for reduce_node in reduce_nodes: - if next_ag.node in ag_nodes_in_path_to_reduce_nodes[reduce_node]: - ag_nodes_in_path_to_reduce_nodes[reduce_node].remove(next_ag.node) - if len(ag_nodes_in_path_to_reduce_nodes[reduce_node]) == 0: - reduces_to_schedule.append(reduce_node) + for scheduled_ag in scheduled_ags: + if scheduled_ag in ag_nodes_in_path_to_reduce_nodes[reduce_node]: + ag_nodes_in_path_to_reduce_nodes[reduce_node].remove(scheduled_ag) + if len(ag_nodes_in_path_to_reduce_nodes[reduce_node]) == 0: + reduces_to_schedule.append(reduce_node) for n in reduces_to_schedule: need_to_schedule = get_node_requirements(n, scheduled) for nn in need_to_schedule: scheduled.append(nn) unscheduled.remove(nn) + if gathered_tracker is not None: + gathered_tracker.apply(nn) # Do the same for output nodes outputs_to_schedule = [] for output_node in output_nodes: - if next_ag.node in ag_nodes_in_path_to_output_nodes[output_node]: - ag_nodes_in_path_to_output_nodes[output_node].remove(next_ag.node) - if len(ag_nodes_in_path_to_output_nodes[output_node]) == 0: - outputs_to_schedule.append(output_node) + for scheduled_ag in scheduled_ags: + if scheduled_ag in ag_nodes_in_path_to_output_nodes[output_node]: + ag_nodes_in_path_to_output_nodes[output_node].remove(scheduled_ag) + if len(ag_nodes_in_path_to_output_nodes[output_node]) == 0: + outputs_to_schedule.append(output_node) for n in outputs_to_schedule: need_to_schedule = get_node_requirements(n, scheduled) for nn in need_to_schedule: scheduled.append(nn) unscheduled.remove(nn) + if gathered_tracker is not None: + gathered_tracker.apply(nn) # print(f"After ag scheduled: scheduled: {scheduled}") @@ -433,9 +815,12 @@ def fast_free_schedule(graph: Graph, available_mem: int, output_size: int, debug continue scheduled.append(node) unscheduled.remove(node) + if gathered_tracker is not None: + gathered_tracker.apply(node) assert len(unscheduled) == 0, f"There are unscheduled nodes: {unscheduled}" ret_graph = make_graph_from_schedule(scheduled) + setattr(ret_graph, SCHEDULER_BUDGET_DIAGNOSTICS_ATTR, diagnostics) ret_graph.lint() return ret_graph diff --git a/deepspeed/compile/passes/prefetch.py b/deepspeed/compile/passes/prefetch.py index 7931bc92d9dd..6bb1c82119e9 100644 --- a/deepspeed/compile/passes/prefetch.py +++ b/deepspeed/compile/passes/prefetch.py @@ -14,6 +14,7 @@ from ..profilers.comm_profile import create_predictor from ..profilers.graph_profile import is_profile_incomplete from ..graph_param import DSGraphParamManager +from ..util import all_reduce from .contract import PassContract, CAP_Z3_GATHER_RELEASE NAME = "prefetch" @@ -43,6 +44,7 @@ def schedule_prefetch(gm: GraphModule, graph_id: int, graph_order: List[Tuple[in bwd: bool) -> GraphModule: profile = profiling_results[graph_id] + process_group = getattr(profile, "process_group", None) profile_graph = profile.bwd_graph if bwd else profile.fwd_graph mem_complete = profile.bwd_mem_complete if bwd else profile.fwd_mem_complete if is_profile_incomplete(profile_graph) or not mem_complete: @@ -51,7 +53,7 @@ def schedule_prefetch(gm: GraphModule, graph_id: int, graph_order: List[Tuple[in max_mem = get_accelerator().total_memory() * (1 - MARGIN) vals_to_bcast = torch.tensor([max_mem], device=torch.device(get_accelerator().current_device())) - dist.all_reduce(vals_to_bcast, dist.ReduceOp.MIN) + all_reduce(vals_to_bcast, dist.ReduceOp.MIN, process_group) max_mem = vals_to_bcast[0].item() mem = profiling_results[graph_id].bwd_mem if bwd else profiling_results[graph_id].fwd_mem diff --git a/deepspeed/compile/passes/selective_gather.py b/deepspeed/compile/passes/selective_gather.py index ddbef1784e9e..45581ea07152 100644 --- a/deepspeed/compile/passes/selective_gather.py +++ b/deepspeed/compile/passes/selective_gather.py @@ -13,7 +13,7 @@ from deepspeed.accelerator import get_accelerator from deepspeed.utils import log_dist -from ..util import get_deepcompile_handle +from ..util import get_deepcompile_handle, all_reduce from ..graph_param import DSGraphParamManager from ..profilers.graph_profile import is_profile_incomplete from .contract import PassContract, CAP_Z3_GATHER_RELEASE @@ -78,6 +78,7 @@ def _profile_result_incomplete(prof) -> bool: def selective_gather(gm: GraphModule, graph_id: int, graph_order: List[Tuple[int, bool]], profiling_results, create_inputs_fn, mem_budget: float, param_manager: DSGraphParamManager, bwd: bool) -> GraphModule: + """Persist high-value parameters only on the final backward graph and within measured headroom.""" target_graph_id = graph_id if not bwd: @@ -93,6 +94,7 @@ def selective_gather(gm: GraphModule, graph_id: int, graph_order: List[Tuple[int if last_backward_graph_id is None or graph_id != last_backward_graph_id: return gm + process_group = getattr(profiling_results[graph_id], "process_group", None) incomplete_profile_ids = [ profile_graph_id for profile_graph_id, prof in profiling_results.items() if _profile_result_incomplete(prof) ] @@ -170,13 +172,16 @@ def selective_gather(gm: GraphModule, graph_id: int, graph_order: List[Tuple[int current_available_mem = accelerator.available_memory() vals_to_bcast = torch.tensor([total_mem, current_available_mem], device=torch.device(get_accelerator().current_device())) - dist.all_reduce(vals_to_bcast, dist.ReduceOp.MIN) + all_reduce(vals_to_bcast, dist.ReduceOp.MIN, process_group) total_mem = vals_to_bcast[0].item() current_available_mem = vals_to_bcast[1].item() budget = _compute_persistence_budget(all_graph_mem_records, total_mem, MEM_MARGIN) profiled_available_mem = budget["available_mem"] - available_mem = profiled_available_mem + current_available_headroom = max(0, int(current_available_mem) - int(total_mem * MEM_MARGIN)) + # The profile models transient peaks, while current allocator state captures + # memory retained since profiling. Persistence must satisfy both views. + available_mem = min(profiled_available_mem, current_available_headroom) ds_id_to_param = {} for g_id, g_pm in param_manager.items(): @@ -190,7 +195,8 @@ def selective_gather(gm: GraphModule, graph_id: int, graph_order: List[Tuple[int f"selective_gather target_graph_id={target_graph_id} profiled_mem_lists={budget['profiled_list_count']} " f"total_mem={total_mem} usable_mem={budget['usable_mem']} peak_resident_alloc={budget['peak_resident_alloc']} " f"transient_peak={budget['transient_peak']} current_available_mem={current_available_mem} " - f"profiled_transient_available_mem={profiled_available_mem} " + f"current_available_headroom={current_available_headroom} " + f"profiled_transient_available_mem={profiled_available_mem} effective_available_mem={available_mem} " f"persistent_count={len(persistent_ds_ids)} persistent_bytes={persistent_bytes} " f"candidate_count={len(ds_ids)} candidate_bytes={candidate_bytes}") @@ -203,7 +209,7 @@ def selective_gather(gm: GraphModule, graph_id: int, graph_order: List[Tuple[int return gm if available_mem == 0: - print_rank_0("selective_gather no profiled headroom for new persistent params") + print_rank_0("selective_gather no effective headroom for new persistent params") return gm persistent_mem = 0 diff --git a/deepspeed/compile/passes/zero3_compile.py b/deepspeed/compile/passes/zero3_compile.py index f39cc3b4b8a3..e0cf6923e37c 100644 --- a/deepspeed/compile/passes/zero3_compile.py +++ b/deepspeed/compile/passes/zero3_compile.py @@ -4,6 +4,9 @@ # DeepSpeed Team import gc +import hashlib +import json +import os from typing import List, Dict, Tuple import _operator @@ -11,21 +14,294 @@ from torch.fx import Graph, Node, GraphModule from ..util import get_input_nodes, get_param_nodes, get_index_by_graph_id, get_deepcompile_handle, get_real_uses, is_cast_op +from ..util import all_reduce, get_rank from ..fx import (add_postprocess, _make_node_meta, get_output_node, move_primals_to_head, add_end_backward, replace_reduce_outputs_with_none, should_release_reduce_buckets) -from ..profilers.graph_profile import ProfilingInterpreter -from ..list_schedule import fast_free_schedule +from ..profilers.graph_profile import ProfilingInterpreter, is_profile_incomplete +from ..list_schedule import (SCHEDULER_BUDGET_DIAGNOSTICS_ATTR, SchedulerMemoryBudget, allgather_allocation_bytes, + fast_free_schedule, max_possible_gathered_bytes, profiled_non_gathered_peak) from .contract import PassContract, CAP_Z3_GATHER_RELEASE import deepspeed.comm as dist from deepspeed.accelerator import get_accelerator +from deepspeed.utils.allocator_telemetry import enabled as allocator_telemetry_enabled +from deepspeed.utils.allocator_telemetry import record_empty_cache, record_scheduler_decision NAME = "zero3_compile" # Inserts the ZeRO-3 all-gather/release ops that prefetch and selective-gather later build on. CONTRACT = PassContract(provides=frozenset({CAP_Z3_GATHER_RELEASE})) +SCHEDULER_DEBUG_ENV = "DEEPSPEED_COMPILE_SCHEDULER_BUDGET_DEBUG" +SCHEDULER_DEBUG_ENV_LEGACY = "DEEPSPEED_DEEPCOMPILE_SCHEDULER_DEBUG" -def add_allgather(graph_id: int, graph: Graph, node: Node, ds_id: int, dtype: torch.dtype): +def _reduce_int(value: int, op, process_group=None): + """Reduce an integer scheduling input, or preserve it outside distributed mode.""" + if not dist.is_initialized(): + return int(value) + + value_tensor = torch.tensor([int(value)], + device=torch.device(get_accelerator().current_device()), + dtype=torch.int64) + all_reduce(value_tensor, op, process_group) + return int(value_tensor.item()) + + +def _rank_min_total_memory(process_group=None): + return _reduce_int(get_accelerator().total_memory(), dist.ReduceOp.MIN, process_group) + + +def _world_size(process_group=None): + if dist.is_initialized(): + if process_group is None: + return dist.get_world_size() + return dist.get_world_size(group=process_group) + return 1 + + +def _sync_profile_complete(profile_complete: bool, process_group=None): + """Require every rank to have a complete profile before using it for scheduling.""" + if not dist.is_initialized(): + return profile_complete + + complete = torch.tensor([1 if profile_complete else 0], + device=torch.device(get_accelerator().current_device()), + dtype=torch.int) + all_reduce(complete, dist.ReduceOp.MIN, process_group) + return bool(complete.item()) + + +def _operator_profile_complete(graph: Graph): + """Require the graph marker plus per-node absolute start and peak memory.""" + return not is_profile_incomplete(graph) and all( + "profile_mem_start" in node.meta and "profile_mem_peak" in node.meta for node in graph.nodes) + + +def _rank_max_operator_profiled_non_gathered_peak(graph: Graph, process_group=None): + """Return the worst absolute peak after removing profiled gather residency.""" + records = [(node.name, int(node.meta.get("profile_mem_start", 0) + or 0), 0, int(node.meta.get("profile_mem_peak", 0) or 0)) for node in graph.nodes] + return _reduce_int(profiled_non_gathered_peak(graph, records), dist.ReduceOp.MAX, process_group) + + +def _build_scheduler_budget_from_operator_profile(graph: Graph, output_size: int = 0, process_group=None): + """Build a budget only when every node has trustworthy operator profile data.""" + if not _operator_profile_complete(graph): + return None + + scheduler_budget = SchedulerMemoryBudget.from_profiled_non_gathered_peak( + _rank_min_total_memory(process_group), _rank_max_operator_profiled_non_gathered_peak(graph, process_group), + output_size) + minimum_budget = _minimum_gather_residency_budget(graph, process_group) + if scheduler_budget is not None and minimum_budget is not None: + scheduler_budget = scheduler_budget.clamped_to_minimum_gather_residency(minimum_budget.max_gathered_bytes) + return scheduler_budget + + +def _minimum_gather_residency_budget(graph: Graph, process_group=None): + """Build a rank-consistent gather-only cap without inferring memory headroom.""" + local_max_single_gather = max((int(node.meta.get("allgather_allocation_bytes", 0) or 0) + for node in graph.nodes if node.target == torch.ops.dc.allgather_param.default), + default=0) + max_single_gather = _reduce_int(local_max_single_gather, dist.ReduceOp.MAX, process_group) + return SchedulerMemoryBudget.minimum_gather_residency(max_single_gather) + + +def _scheduler_debug_enabled(): + return any( + os.environ.get(env_name, "").lower() not in ("", "0", "false", "no") + for env_name in (SCHEDULER_DEBUG_ENV, SCHEDULER_DEBUG_ENV_LEGACY)) + + +def _print_scheduler_debug(message: str, process_group=None): + if not _scheduler_debug_enabled(): + return + if not dist.is_initialized() or get_rank(process_group) == 0: + print(message, flush=True) + + +def _stable_schedule_target(target): + if isinstance(target, str): + return target + module = getattr(target, "__module__", None) + qualname = getattr(target, "__qualname__", None) or getattr(target, "__name__", None) + if module and qualname: + return f"{module}.{qualname}" + return f"{type(target).__module__}.{type(target).__qualname__}" + + +def _final_schedule_fingerprint(graph: Graph): + """Hash stable node order and dependencies without process-local graph identifiers.""" + entries = [] + for node in graph.nodes: + inputs = ",".join(input_node.name for input_node in node.all_input_nodes) + entries.append(f"{node.op}|{node.name}|{_stable_schedule_target(node.target)}|{inputs}") + digest = hashlib.sha256("\n".join(entries).encode("utf-8")).digest() + return int.from_bytes(digest[:8], byteorder="big") & ((1 << 63) - 1) + + +def _collective_schedule_projection(graph: Graph): + """Return the ordered gather/release identities, excluding graph-local names and IDs.""" + entries = [] + for node in graph.nodes: + if node.target == torch.ops.dc.allgather_param.default: + kind = "gather" + elif node.target == torch.ops.dc.release_param.default: + kind = "release" + else: + continue + if len(node.args) < 3 or not isinstance(node.args[2], int): + raise RuntimeError(f"DeepCompile ZeRO-3 {kind} node has no stable integer ds_id: {node.format_node()}") + entries.append((kind, node.args[2])) + return entries + + +def _collective_schedule_fingerprint(projection): + payload = json.dumps(projection, separators=(",", ":")) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def _validate_final_schedule_fingerprint(graph: Graph, graph_id: int, bwd: bool, process_group=None): + """In scheduler debug mode, fail every rank when final graph order diverges.""" + if not _scheduler_debug_enabled(): + return None + + fingerprint = _final_schedule_fingerprint(graph) + collective_projection = _collective_schedule_projection(graph) + collective_fingerprint = _collective_schedule_fingerprint(collective_projection) + collective_message = ("DeepCompile ZeRO-3 collective_schedule_projection " + f"graph_id={graph_id} bwd={bwd} value={collective_fingerprint} " + f"entries={len(collective_projection)} sequence=" + f"{json.dumps(collective_projection, separators=(',', ':'))}") + if not dist.is_initialized(): + _print_scheduler_debug(collective_message, process_group) + _print_scheduler_debug( + f"DeepCompile ZeRO-3 final_schedule_fingerprint graph_id={graph_id} bwd={bwd} value={fingerprint}", + process_group) + return fingerprint + + _print_scheduler_debug(collective_message, process_group) + device = torch.device(get_accelerator().current_device()) + min_fingerprint = torch.tensor([fingerprint], device=device, dtype=torch.int64) + max_fingerprint = min_fingerprint.clone() + all_reduce(min_fingerprint, dist.ReduceOp.MIN, process_group) + all_reduce(max_fingerprint, dist.ReduceOp.MAX, process_group) + if min_fingerprint.item() != max_fingerprint.item(): + raise RuntimeError( + f"DeepCompile ZeRO-3 final schedule fingerprint mismatch for graph_id={graph_id} bwd={bwd}: " + f"min={min_fingerprint.item()} max={max_fingerprint.item()}") + _print_scheduler_debug( + f"DeepCompile ZeRO-3 final_schedule_fingerprint graph_id={graph_id} bwd={bwd} value={fingerprint}", + process_group) + return fingerprint + + +def _set_allgather_allocation_metadata(graph: Graph, process_group=None): + """Stamp padded gather allocation bytes without discarding a more precise estimate.""" + world_size = None + for node in graph.nodes: + if node.target == torch.ops.dc.allgather_param.default: + if world_size is None: + world_size = _world_size(process_group) + dtype = node.kwargs.get("dtype") if isinstance(node.kwargs, dict) else None + profiled_bytes = allgather_allocation_bytes(node.meta.get("tensor_size", 0), dtype, world_size) + node.meta["allgather_allocation_bytes"] = max(int(node.meta.get("allgather_allocation_bytes", 0) or 0), + profiled_bytes) + + +def _scheduler_budget_disabled_reason(graph: Graph, scheduler_budget): + if scheduler_budget is not None: + return None + if not _operator_profile_complete(graph): + return "incomplete_operator_profile" + return "invalid_profiled_non_gathered_peak" + + +def _scheduler_budget_from_operator_profile(gm: GraphModule, process_group=None): + """Derive a rank-consistent budget and explain why a non-constraining gate is disabled.""" + if not dist.is_initialized(): + return None, "non_distributed" + + _set_allgather_allocation_metadata(gm.graph, process_group) + operator_profile_complete = _sync_profile_complete(_operator_profile_complete(gm.graph), process_group) + if not operator_profile_complete: + # An unvisited suffix can exceed every observed partial peak, so do + # not infer whole-graph headroom from partial data or allocator state. + # Still constrain scheduler-controlled gathered-parameter residency to + # one largest-gather-sized cap so incomplete profiling cannot silently + # restore unconstrained gather overlap. + scheduler_budget = _minimum_gather_residency_budget(gm.graph, process_group) + if scheduler_budget is None: + return None, "incomplete_operator_profile_without_gather" + return scheduler_budget, None + + scheduler_budget = _build_scheduler_budget_from_operator_profile(gm.graph, process_group=process_group) + # A gate larger than every gather combined cannot affect ordering, so keep + # legacy behavior and make the disabled state explicit in diagnostics. + if scheduler_budget is not None and scheduler_budget.max_gathered_bytes >= max_possible_gathered_bytes(gm.graph): + return None, "budget_not_constraining" + return scheduler_budget, _scheduler_budget_disabled_reason(gm.graph, scheduler_budget) + + +def _log_scheduler_result(graph_id: int, + bwd: bool, + scheduler_budget, + disabled_reason, + graph: Graph, + process_group=None): + diagnostics = getattr(graph, SCHEDULER_BUDGET_DIAGNOSTICS_ATTR, {}) + selected = diagnostics.get("selected", []) + max_live_gathered_bytes = max((entry.get("peak_gathered_bytes", 0) for entry in selected), default=0) + rejected_candidates = diagnostics.get("budget_rejected_candidates", []) + minimum_rejected_candidate_peak_bytes = None + if (scheduler_budget is not None and rejected_candidates + and (_scheduler_debug_enabled() or allocator_telemetry_enabled())): + minimum_rejected_candidate_peak_bytes = min( + scheduler_budget.max_gathered_bytes + entry.get("over_budget_bytes", 0) for entry in rejected_candidates) + record_scheduler_decision( + graph_id=graph_id, + backward=bwd, + budget_source=scheduler_budget.source if scheduler_budget is not None else None, + disabled_reason=disabled_reason, + max_gathered_bytes=scheduler_budget.max_gathered_bytes if scheduler_budget is not None else None, + max_live_gathered_bytes=max_live_gathered_bytes, + budget_rejections=diagnostics.get("budget_rejections", 0), + over_budget_fallbacks=len(diagnostics.get("budget_overflows", [])), + minimum_rejected_candidate_peak_bytes=minimum_rejected_candidate_peak_bytes, + ) + if scheduler_budget is None: + _print_scheduler_debug( + f"DeepCompile ZeRO-3 scheduler graph_id={graph_id} bwd={bwd} budget_enabled=False " + f"disabled_reason={disabled_reason} selected_count={len(selected)} " + f"max_live_gathered_bytes={max_live_gathered_bytes}", process_group) + return + + _print_scheduler_debug( + f"DeepCompile ZeRO-3 scheduler graph_id={graph_id} bwd={bwd} budget_enabled=True " + f"budget_source={scheduler_budget.source} max_gathered_bytes={scheduler_budget.max_gathered_bytes} " + f"safety_margin={scheduler_budget.safety_margin} " + f"profiled_non_gathered_peak_mem={scheduler_budget.profiled_non_gathered_peak_mem} " + f"budget_rejections={diagnostics.get('budget_rejections', 0)} " + f"minimum_rejected_candidate_peak_bytes={minimum_rejected_candidate_peak_bytes} " + f"over_budget_fallbacks={len(diagnostics.get('budget_overflows', []))} " + f"max_live_gathered_bytes={max_live_gathered_bytes}", process_group) + + +def _dtype_element_size(dtype: torch.dtype): + return torch.empty((), dtype=dtype).element_size() + + +def _param_allgather_allocation_bytes(param, dtype: torch.dtype): + """Return the registered parameter size in its target gather dtype.""" + return int(param.numel) * _dtype_element_size(dtype) + + +def add_allgather(graph_id: int, + graph: Graph, + node: Node, + ds_id: int, + dtype: torch.dtype, + allgather_allocation_bytes: int = None): + """Insert gather and wait nodes while preserving the original graph output edge.""" new_ag_node = add_postprocess(graph, node, torch.ops.dc.allgather_param.default, @@ -33,6 +309,8 @@ def add_allgather(graph_id: int, graph: Graph, node: Node, ds_id: int, dtype: to extra_kwargs={"dtype": dtype}, name=f"allgather_ds_param_{node.target}_{ds_id}", meta=_make_node_meta(node, ds_id, True)) + if allgather_allocation_bytes is not None: + new_ag_node.meta["allgather_allocation_bytes"] = int(allgather_allocation_bytes) new_ag_node.meta["val"] = node.meta["val"].to(dtype) # Set the previous node back to output @@ -73,6 +351,7 @@ def add_reduce(graph_id: int, graph: Graph, grad_node: Node, param_name: str, ds def add_gather_and_release(graph_id: int, graph: Graph, param_manager, param_nodes: List[Node]) -> Graph: + """Insert gather/wait lifetimes and attach releases to ordinary parameter consumers.""" node_to_uses = get_real_uses(graph) for pn in param_nodes: @@ -90,7 +369,14 @@ def add_gather_and_release(graph_id: int, graph: Graph, param_manager, param_nod fuse_typecast = True target_dtype = casted_dtype - add_allgather(graph_id, graph, pn, param_manager.ds_ids[pn.name], target_dtype) + param = param_manager.params[pn.name] + allgather_node = add_allgather(graph_id, + graph, + pn, + param_manager.ds_ids[pn.name], + target_dtype, + allgather_allocation_bytes=_param_allgather_allocation_bytes( + param, target_dtype)) if fuse_typecast: users = node_to_uses[typecast_node] wait_node = typecast_node.args[0] @@ -101,6 +387,15 @@ def add_gather_and_release(graph_id: int, graph: Graph, param_manager, param_nod graph.erase_node(typecast_node) else: users = node_to_uses[pn] + if len(users) == 0: + # Parameters returned directly by the graph have no ordinary + # consumer to trigger gathering, so make the waited gather the + # output while retaining its original output name. + output_node = get_output_node(graph) + wait_node = next(user for user in allgather_node.users + if user.target == torch.ops.dc.wait_allgather.default) + wait_node.meta["original_output_name"] = pn.name + output_node.replace_input_with(pn, wait_node) ds_id = param_manager.ds_ids[pn.name] for user in users: @@ -123,6 +418,7 @@ def add_gather_and_release(graph_id: int, graph: Graph, param_manager, param_nod def add_gather_and_reduce(graph_id: int, graph: Graph, param_manager, param_nodes_bw: List[Node], param_name_to_grad: Dict[str, Node]) -> Graph: + """Add parameter lifetimes and gradient reductions to a backward graph.""" add_gather_and_release(graph_id, graph, param_manager, param_nodes_bw) @@ -141,24 +437,29 @@ def add_z3_gather_release_fw(gm: GraphModule, create_inputs_fn, param_manager, debug_log=False) -> GraphModule: + """Profile, budget, and schedule ZeRO-3 parameter lifetimes for a forward graph.""" nz3 = get_deepcompile_handle() real_inputs = create_inputs_fn() param_indices = profiling_results[graph_id].param_indices + process_group = getattr(profiling_results[graph_id], "process_group", None) gm.graph = add_gather_and_release(graph_id, gm.graph, param_manager[graph_id], get_param_nodes(gm.graph, param_indices)) nz3.register_graph_z3(graph_id, [v[1] for v in param_indices]) # Need this before profiling - profiler = ProfilingInterpreter(gm, debug_log=debug_log) + profiler = ProfilingInterpreter(gm, debug_log=debug_log, process_group=process_group) profiler.run(*real_inputs) del profiler gc.collect() - get_accelerator().empty_cache() + record_empty_cache("zero3-compile.forward-post-profile", get_accelerator().empty_cache) + # Build the shared scheduling budget after the operator profile is complete + # but before the scheduler rewrites graph order and Inductor metadata. + scheduler_budget, disabled_reason = _scheduler_budget_from_operator_profile(gm, process_group) - rank = dist.get_rank() + rank = get_rank(process_group) graph_index = get_index_by_graph_id(graph_order, graph_id) if rank == 0 and debug_log: print(f"Fwd before scheduling graph {graph_index} graph_id={graph_id} {gm.graph}") @@ -173,7 +474,15 @@ def add_z3_gather_release_fw(gm: GraphModule, gm.graph, get_accelerator().available_memory(), 0, # unused - debug_log=debug_log) + debug_log=debug_log, + scheduler_budget=scheduler_budget) + _log_scheduler_result(graph_id, + bwd=False, + scheduler_budget=scheduler_budget, + disabled_reason=disabled_reason, + graph=gm.graph, + process_group=process_group) + _validate_final_schedule_fingerprint(gm.graph, graph_id, bwd=False, process_group=process_group) if rank == 0 and debug_log: print(f"Fwd after scheduling graph {graph_index} graph_id={graph_id} {gm.graph}") @@ -188,21 +497,26 @@ def add_z3_gather_release_bw(gm: GraphModule, create_inputs_fn, param_manager, debug_log=False) -> GraphModule: + """Profile, budget, and schedule gathers, releases, and reductions for backward.""" param_nodes_bw, param_name_to_grad = param_manager[graph_id].get_bwd_mapping(gm.graph) gm.graph = add_gather_and_reduce(graph_id, gm.graph, param_manager[graph_id], param_nodes_bw, param_name_to_grad) input_nodes = get_input_nodes(gm.graph) real_inputs = create_inputs_fn() + process_group = getattr(profiling_results[graph_id], "process_group", None) assert len(input_nodes) == len(real_inputs), f"Expected {len(real_inputs)} inputs, got {len(input_nodes)}" - real_outputs = ProfilingInterpreter(gm, debug_log=debug_log).run(*real_inputs) + real_outputs = ProfilingInterpreter(gm, debug_log=debug_log, process_group=process_group).run(*real_inputs) del real_outputs gc.collect() - get_accelerator().empty_cache() + record_empty_cache("zero3-compile.backward-post-profile", get_accelerator().empty_cache) + # The scheduler consumes only DP-group-reduced inputs, ensuring every group + # rank emits collectives in the same order even when allocator state differs. + scheduler_budget, disabled_reason = _scheduler_budget_from_operator_profile(gm, process_group) - rank = dist.get_rank() + rank = get_rank(process_group) graph_index = get_index_by_graph_id(graph_order, graph_id) if rank == 0 and debug_log: print(f"Bwd before scheduling graph {graph_index} graph_id={graph_id} {gm.graph}") @@ -211,10 +525,18 @@ def add_z3_gather_release_bw(gm: GraphModule, gm.graph, get_accelerator().available_memory(), 0, # unused - debug_log=debug_log) + debug_log=debug_log, + scheduler_budget=scheduler_budget) + _log_scheduler_result(graph_id, + bwd=True, + scheduler_budget=scheduler_budget, + disabled_reason=disabled_reason, + graph=gm.graph, + process_group=process_group) add_end_backward(gm.graph, graph_id, should_release_reduce_buckets(graph_order, graph_id)) replace_reduce_outputs_with_none(gm.graph) + _validate_final_schedule_fingerprint(gm.graph, graph_id, bwd=True, process_group=process_group) return gm diff --git a/deepspeed/compile/patch_compiled_func.py b/deepspeed/compile/patch_compiled_func.py index c77d529a64ac..2a0ffd593129 100644 --- a/deepspeed/compile/patch_compiled_func.py +++ b/deepspeed/compile/patch_compiled_func.py @@ -82,12 +82,21 @@ class PatchedFunction(torch.autograd.Function, metaclass=FunctionMeta): def unpatch_compiled_func(): + """Restore torch.autograd.Function and discard inputs captured for this compile cycle.""" global enabled_patched_func enabled_patched_func = False global original_grad_fn - torch.autograd.Function = original_grad_fn + if original_grad_fn is not None: + torch.autograd.Function = original_grad_fn + original_grad_fn = None + clear_backward_inputs() def get_backward_inputs(): return backward_inputs + + +def clear_backward_inputs(): + """Drop captured real backward inputs before the next graph compilation.""" + backward_inputs.clear() diff --git a/deepspeed/compile/profilers/__init__.py b/deepspeed/compile/profilers/__init__.py index 6134504cc4dc..eddd48a3e59f 100644 --- a/deepspeed/compile/profilers/__init__.py +++ b/deepspeed/compile/profilers/__init__.py @@ -3,7 +3,7 @@ # DeepSpeed Team -from typing import List, Tuple +from typing import Any, List, Tuple from dataclasses import dataclass, field from torch.fx import Graph @@ -23,3 +23,6 @@ class ProfilingResult: fwd_tensor_sizes: List[Tuple[str, int]] = field(default_factory=list) # name, size bwd_tensor_sizes: List[Tuple[str, int]] = field(default_factory=list) param_indices: List[Tuple[int, int, Tuple[int, ...]]] = field(default_factory=list) # index, ds_id, ds_shape + # Keep newly added fields at the end so positional construction of the + # long-standing profiling fields remains backward compatible. + process_group: Any = None diff --git a/deepspeed/compile/profilers/graph_profile.py b/deepspeed/compile/profilers/graph_profile.py index ec26e029cbb0..f44efd07f2f5 100644 --- a/deepspeed/compile/profilers/graph_profile.py +++ b/deepspeed/compile/profilers/graph_profile.py @@ -20,7 +20,7 @@ import deepspeed.comm as dist from deepspeed.accelerator import get_accelerator -from ..util import is_comm_op, is_release_node, get_deepcompile_handle +from ..util import is_comm_op, is_release_node, get_deepcompile_handle, all_reduce, get_rank, barrier def _all_real_if_tensor(args): @@ -60,6 +60,8 @@ def _node_size(out): "tensor_size": 0, "alloc_mem": 0, "max_mem": 0, + "profile_mem_start": 0, + "profile_mem_peak": 0, } _PROFILE_INCOMPLETE_ATTR = "_deepcompile_profile_incomplete" _PROFILE_INCOMPLETE_META_KEY = "deepcompile_profile_incomplete" @@ -91,6 +93,14 @@ def _backfill_missing_profile_metadata(graph: Graph, profile_complete: bool = Tr node.meta.setdefault(key, default) +def _clear_interpreter_env(interpreter: Interpreter): + """Release FX interpreter references so profiling outputs do not remain live.""" + try: + interpreter.env.clear() + except Exception: + pass + + def _run_warmup_for_profile(call_fn, warmup): for _ in range(warmup): warmup_out = call_fn() @@ -117,12 +127,17 @@ def _get_mem_usage_out_of_torch(): import pynvml pynvml.nvmlInit() - current_dev_id = get_accelerator().current_device() - handle = pynvml.nvmlDeviceGetHandleByIndex(current_dev_id) + accelerator = get_accelerator() + current_dev_id = accelerator.current_device() + map_nvml_device = getattr(accelerator, "_get_nvml_gpu_id", None) + nvml_dev_id = map_nvml_device(current_dev_id) if callable(map_nvml_device) else current_dev_id + handle = pynvml.nvmlDeviceGetHandleByIndex(nvml_dev_id) info = pynvml.nvmlDeviceGetMemoryInfo(handle) - torch_alloc = get_accelerator().memory_allocated() - adjust = info.used - torch_alloc + # NVML includes PyTorch's cached allocator reservation. Subtract the + # whole reservation so later reuse is counted only as live allocation. + torch_reserved = accelerator.memory_reserved() + adjust = max(0, int(info.used) - int(torch_reserved or 0)) except Exception: # pynvml not available pass @@ -130,10 +145,24 @@ def _get_mem_usage_out_of_torch(): return adjust +def _absolute_profile_memory(mem_usage_out_of_torch): + """Read absolute allocator residency and peak with external memory included once.""" + return (int(get_accelerator().memory_allocated()) + int(mem_usage_out_of_torch), + int(get_accelerator().max_memory_allocated()) + int(mem_usage_out_of_torch)) + + +def _rank_max_profile_memory(start_mem, peak_mem, device, distributed, process_group=None): + """Return per-field worst-rank absolute memory without averaging rank asymmetry.""" + values = torch.tensor([int(start_mem), int(peak_mem)], device=device, dtype=torch.int64) + if distributed: + all_reduce(values, dist.ReduceOp.MAX, process_group) + return int(values[0].item()), int(values[1].item()) + + # https://pytorch.org/tutorials/intermediate/fx_profiling_tutorial.html class ProfilingInterpreter(Interpreter): - def __init__(self, gm: GraphModule, iteration: int = 10, warmup: int = 5, debug_log=False): + def __init__(self, gm: GraphModule, iteration: int = 10, warmup: int = 5, debug_log=False, process_group=None): super().__init__(gm) self.nz3 = get_deepcompile_handle() @@ -145,6 +174,7 @@ def __init__(self, gm: GraphModule, iteration: int = 10, warmup: int = 5, debug_ self.device = torch.device(get_accelerator().current_device()) self.cache: Dict[Tuple, Any] = {} self.distributed = dist.is_initialized() + self.process_group = process_group self.allgather_mem: Dict[int, int] = {} self.debug_log = debug_log self.mem_usage_out_of_torch = 0 @@ -168,27 +198,36 @@ def run(self, *args) -> Any: except Exception as e: profile_complete = False msg = e.msg if "msg" in dir(e) else str(e) - if not self.distributed or dist.get_rank() == 0: + if not self.distributed or get_rank(self.process_group) == 0: print(f"DeepCompile profiling failed; using default profile metadata for incomplete nodes: {msg}") finally: + # Keep this try/finally so profiling state is restored if gathered-param cleanup fails. try: self.nz3.clear_all_gathered_params() finally: - try: - self.nz3.enable_profiling(False) - finally: - _backfill_missing_profile_metadata(self.graph, profile_complete=profile_complete) + self.nz3.enable_profiling(False) + _clear_interpreter_env(self) + _backfill_missing_profile_metadata(self.graph, profile_complete=profile_complete) return return_val def run_node(self, n: torch.fx.Node) -> Any: if n.op in {"placeholder", "output"}: + get_accelerator().reset_peak_memory_stats() + profile_mem_start, _ = _absolute_profile_memory(self.mem_usage_out_of_torch) + ret = super().run_node(n) + _, profile_mem_peak = _absolute_profile_memory(self.mem_usage_out_of_torch) + profile_mem_start, profile_mem_peak = _rank_max_profile_memory(profile_mem_start, profile_mem_peak, + self.device, self.distributed, + self.process_group) n.meta["device_time"] = 0.0 n.meta["wall_time"] = 0.0 n.meta["alloc_mem"] = 0 n.meta["max_mem"] = 0 n.meta["tensor_size"] = _node_size(n) - return super().run_node(n) + n.meta["profile_mem_start"] = profile_mem_start + n.meta["profile_mem_peak"] = profile_mem_peak + return ret args, kwargs = self.fetch_args_kwargs_from_env(n) assert isinstance(args, tuple) @@ -215,7 +254,7 @@ def rebuild_param_if_necessary(v): cache_hit_flag = torch.tensor([0 if cache_hit else 1], device=self.device, dtype=torch.int) if self.distributed: - dist.all_reduce(cache_hit_flag, dist.ReduceOp.SUM) + all_reduce(cache_hit_flag, dist.ReduceOp.SUM, self.process_group) cache_hit = cache_hit_flag.item() == 0 if cache_hit: @@ -236,6 +275,7 @@ def rebuild_param_if_necessary(v): get_accelerator().reset_peak_memory_stats() alloc_mem_start = get_accelerator().memory_allocated() max_mem_start = get_accelerator().max_memory_allocated() + profile_mem_start, _ = _absolute_profile_memory(self.mem_usage_out_of_torch) def run_target(): return getattr(self, n.op)(n.target, args, kwargs) @@ -245,7 +285,7 @@ def run_target(): if is_comm_op(n): assert self.distributed, f"Distributed environment is not initialized but comm operator {n.name} {n.target} is used." - dist.barrier() + barrier(self.process_group) start = time.time() out = _run_repeatedly_for_profile(run_target, iteration, start_events, end_events) @@ -253,10 +293,16 @@ def run_target(): walltime_sum = time.time() - start if is_comm_op(n): - dist.barrier() + barrier(self.process_group) alloc_mem = get_accelerator().memory_allocated() - alloc_mem_start + self.mem_usage_out_of_torch max_memory = get_accelerator().max_memory_allocated() - max_mem_start + self.mem_usage_out_of_torch + _, profile_mem_peak = _absolute_profile_memory(self.mem_usage_out_of_torch) + profile_mem_start, profile_mem_peak = _rank_max_profile_memory(profile_mem_start, profile_mem_peak, + self.device, self.distributed, + self.process_group) + n.meta["profile_mem_start"] = profile_mem_start + n.meta["profile_mem_peak"] = profile_mem_peak tensor_size = _node_size(out) def partition_param_if_necessary(v): @@ -276,7 +322,7 @@ def partition_param_if_necessary(v): vals_to_bcast = torch.tensor([device_time, wall_time, alloc_mem, max_memory, tensor_size], device=self.device) if self.distributed: - dist.all_reduce(vals_to_bcast, dist.ReduceOp.AVG) + all_reduce(vals_to_bcast, dist.ReduceOp.AVG, self.process_group) n.meta["device_time"] = vals_to_bcast[0].item() n.meta["wall_time"] = vals_to_bcast[1].item() n.meta["alloc_mem"] = int(vals_to_bcast[2].item()) @@ -288,7 +334,7 @@ def partition_param_if_necessary(v): if is_release_op: n.meta["alloc_mem"] = -self.allgather_mem.get(args[2], 0) - if dist.get_rank() == 0 and self.debug_log: + if get_rank(self.process_group) == 0 and self.debug_log: print( f"{n.target} {n.meta['device_time']:.2f}ms {n.meta['wall_time']:.2f}ms alloc_mem={n.meta['alloc_mem'] / 1024 / 1024:.2f}MB max_mem={n.meta['max_mem'] / 1024 / 1024:.2f}MB tensor_size={n.meta['tensor_size']}" ) @@ -307,25 +353,28 @@ def partition_param_if_necessary(v): class MemoryProfilingInterpreter(Interpreter): - def __init__(self, gm: GraphModule, debug_log=False): + def __init__(self, gm: GraphModule, debug_log=False, process_group=None): super().__init__(gm) self.nz3 = get_deepcompile_handle() self.device = torch.device(get_accelerator().current_device()) self.mem_record = [] self.last_alloc = get_accelerator().memory_allocated() self.profile_complete = True + self.process_group = process_group self.node_counter = 0 self.node_num = len(gm.graph.nodes) self.debug_log = debug_log def run(self, *args) -> Any: + """Profile absolute memory and release gathered/interpreter state on every exit.""" return_val = None self.profile_complete = True try: assert _all_real_if_tensor(args), "Inputs must be real tensors" self.nz3.enable_profiling(True) self.mem_usage_out_of_torch = _get_mem_usage_out_of_torch() + self.last_alloc = int(get_accelerator().memory_allocated()) + int(self.mem_usage_out_of_torch) with unset_fake_temporarily(): with get_accelerator().random().fork_rng(devices=[self.device]): @@ -333,17 +382,21 @@ def run(self, *args) -> Any: except Exception as e: self.profile_complete = False self.mem_record.clear() + _backfill_missing_profile_metadata(self.graph, profile_complete=False) print(f"MemoryProfiling error {e}") finally: + # Keep this try/finally so profiling state is restored if gathered-param cleanup fails. try: self.nz3.clear_all_gathered_params() finally: self.nz3.enable_profiling(False) + _clear_interpreter_env(self) return return_val def run_node(self, n: torch.fx.Node) -> Any: get_accelerator().reset_peak_memory_stats() + profile_mem_start, _ = _absolute_profile_memory(self.mem_usage_out_of_torch) if n.op in {"placeholder", "output"}: ret = super().run_node(n) @@ -355,17 +408,20 @@ def run_node(self, n: torch.fx.Node) -> Any: del args, kwargs - current_alloc = get_accelerator().memory_allocated() + self.mem_usage_out_of_torch - max_alloc = get_accelerator().max_memory_allocated() + self.mem_usage_out_of_torch - vals_to_bcast = torch.tensor([current_alloc, max_alloc], device=self.device, dtype=torch.int64) - dist.all_reduce(vals_to_bcast, dist.ReduceOp.MAX) - current_alloc = vals_to_bcast[0].item() - max_alloc = vals_to_bcast[1].item() + current_alloc, max_alloc = _absolute_profile_memory(self.mem_usage_out_of_torch) + absolute_record = torch.tensor([profile_mem_start, current_alloc, max_alloc], + device=self.device, + dtype=torch.int64) + if dist.is_initialized(): + all_reduce(absolute_record, dist.ReduceOp.MAX, self.process_group) + profile_mem_start, current_alloc, max_alloc = (int(value.item()) for value in absolute_record) + n.meta["profile_mem_start"] = profile_mem_start + n.meta["profile_mem_peak"] = max_alloc self.mem_record.append((n.name, current_alloc, current_alloc - self.last_alloc, max_alloc)) self.node_counter += 1 - if self.debug_log and dist.get_rank() == 0: + if self.debug_log and get_rank(self.process_group) == 0: print( f"Mem prof Node {self.node_counter}/{self.node_num} {n.name} memory {current_alloc / 1024 / 1024:.2f}MB delta {(current_alloc - self.last_alloc) / 1024 / 1024:.2f}MB" ) diff --git a/deepspeed/compile/util.py b/deepspeed/compile/util.py index f771d8051127..f300c42cfd15 100644 --- a/deepspeed/compile/util.py +++ b/deepspeed/compile/util.py @@ -90,6 +90,24 @@ def log_rank0(msg: str, enable: bool = False): print(msg) +def all_reduce(tensor, op, process_group=None): + if process_group is None: + return dist.all_reduce(tensor, op) + return dist.all_reduce(tensor, op, group=process_group) + + +def get_rank(process_group=None): + if process_group is None: + return dist.get_rank() + return dist.get_rank(group=process_group) + + +def barrier(process_group=None): + if process_group is None: + return dist.barrier() + return dist.barrier(group=process_group) + + @functools.lru_cache def get_no_copy_ops(): # Need to compile custom ops @@ -251,6 +269,7 @@ def convert(t): def get_last_uses(graph: Graph): + """Map values to last consumers while propagating lifetimes through no-copy ops.""" position = {node: i for i, node in enumerate(graph.nodes)} node_to_last_use: Dict[Node, Node] = {} @@ -262,7 +281,9 @@ def register_last_uses(n: Node, user: Node): known_last_use = None if user.target in no_copy_ops and n in node_to_last_use: - last_user = node_to_last_use[user] + # A no-copy node can itself be user-less (for example, a graph + # output alias). In that case its own position is the lifetime end. + last_user = node_to_last_use.get(user, user) last_use_position = position[last_user] known_last_use = node_to_last_use[n] @@ -271,7 +292,7 @@ def register_last_uses(n: Node, user: Node): if n not in node_to_last_use or update: if user.target in no_copy_ops: - user = node_to_last_use[user] + user = node_to_last_use.get(user, user) node_to_last_use[n] = user user_to_last_uses.setdefault(user, []).append(n) diff --git a/deepspeed/compile/z3_eager_fallback.py b/deepspeed/compile/z3_eager_fallback.py index 27450c086e77..e6512ade5f2f 100644 --- a/deepspeed/compile/z3_eager_fallback.py +++ b/deepspeed/compile/z3_eager_fallback.py @@ -4,6 +4,7 @@ # DeepSpeed Team from contextlib import contextmanager +import sys import torch @@ -17,6 +18,16 @@ def get_active_z3_eager_fallback(): return _ACTIVE_FALLBACK +def is_dynamo_guard_evaluation(): + """Return whether the current parameter access originates from Dynamo guard evaluation.""" + frame = sys._getframe() + while frame is not None: + if frame.f_globals.get("__name__") == "torch._dynamo.guards": + return True + frame = frame.f_back + return False + + def record_z3_eager_fallback_param(param): fallback = get_active_z3_eager_fallback() if fallback is None: diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 194dc3555387..f1ddf687ed3d 100755 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -136,6 +136,7 @@ from deepspeed.profiling.flops_profiler.profiler import FlopsProfiler from deepspeed.utils.logging import print_json_dist, print_configuration, set_log_level_from_string +from deepspeed.utils.allocator_telemetry import record_empty_cache from deepspeed.accelerator import get_accelerator @@ -772,11 +773,18 @@ def __del__(self): logger.debug("DeepSpeedEngine.__del__ cleanup skipped: %s", exc, exc_info=True) def destroy(self): + self._release_deepcompile_compiled_backward_state() + self._release_deepcompile_dynamo_config() optimizer = getattr(self, "optimizer", None) if optimizer is not None and hasattr(optimizer, 'destroy'): optimizer.destroy() if self.is_deepcompile_active(): - get_deepcompile_handle().cleanup() + try: + get_deepcompile_handle().cleanup() + finally: + # Native cleanup is process-global and must run only once even + # when destroy() is followed by __del__(). + self._set_deepcompile_active(False) debug_clear_module_and_param_names() checkpoint_engine = getattr(self, "checkpoint_engine", None) @@ -5410,7 +5418,7 @@ def empty_partition_cache(self): if hasattr(self.optimizer, 'empty_partition_cache'): self.optimizer.empty_partition_cache() gc.collect() - get_accelerator().empty_cache() + record_empty_cache("engine.empty-partition-cache", get_accelerator().empty_cache) def get_autosp_backend(self, compile_kwargs): if self.compile_autosp() and self.zero_optimization_stage() not in [ @@ -5523,6 +5531,10 @@ def compile(self, def _set_deepcompile_active(self, active: bool) -> None: """Toggle DeepCompile runtime state and manage forward hooks accordingly.""" + if not active: + self._release_deepcompile_compiled_backward_state() + self._release_deepcompile_dynamo_config() + if self._deepcompile_active == active: return @@ -5541,6 +5553,18 @@ def _set_deepcompile_active(self, active: bool) -> None: self._deepcompile_active = active + def _release_deepcompile_compiled_backward_state(self) -> None: + owned_frames = getattr(self, "_deepcompile_owned_frames", None) + if owned_frames: + from deepspeed.compile.backend import cleanup_compiled_backward_state + cleanup_compiled_backward_state(owned_frames=owned_frames) + + def _release_deepcompile_dynamo_config(self) -> None: + restore_dynamo_config = getattr(self, "_deepcompile_dynamo_config_restore", None) + if restore_dynamo_config is not None: + restore_dynamo_config() + del self._deepcompile_dynamo_config_restore + def get_compile_time(self): from deepspeed.compile.backend import opt_pass_times return opt_pass_times diff --git a/deepspeed/runtime/zero/parameter_offload.py b/deepspeed/runtime/zero/parameter_offload.py index 161b3e27e440..7d70efc225dc 100644 --- a/deepspeed/runtime/zero/parameter_offload.py +++ b/deepspeed/runtime/zero/parameter_offload.py @@ -63,8 +63,12 @@ def __getitem__(self, key): if hasattr(param, "ds_status") and param.ds_status == ZeroParamStatus.NOT_AVAILABLE: if self._parent_module._parameters._in_forward and not torch.compiler.is_compiling(): - from deepspeed.compile.z3_eager_fallback import get_active_z3_eager_fallback + from deepspeed.compile.z3_eager_fallback import get_active_z3_eager_fallback, is_dynamo_guard_evaluation fallback = get_active_z3_eager_fallback() + if fallback is not None and is_dynamo_guard_evaluation(): + # Dynamo guards only inspect parameter identity and metadata. + # Gathering here retains a full parameter before the compiled forward. + return param if fallback is None: register_external_parameter(FWD_MODULE_STACK[-1], param) param.all_gather() diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index 4a66b85e8884..9e094c2f289d 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -17,6 +17,7 @@ from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors from deepspeed.runtime.base_optimizer import ZeROOptimizer from deepspeed.utils import logger +from deepspeed.utils.allocator_telemetry import record_allocator_retry_sample, record_empty_cache from deepspeed.utils.torch import register_grad_hook, required_torch_version from deepspeed.runtime.fp16.loss_scaler import CreateLossScaler from deepspeed.runtime.torch_autocast import get_autocast_dtype, get_all_comm_dtypes, is_autocast_initialized, sort_dtypes @@ -2592,6 +2593,7 @@ def step(self, closure=None): alloc_retries = memory_stats.get("num_alloc_retries") if alloc_retries is None: alloc_retries = 0 + record_allocator_retry_sample(self.n_caching_allocator_flushes, alloc_retries) if alloc_retries > self.n_caching_allocator_flushes: if dist.get_rank() == 0: logger.warning( @@ -3587,7 +3589,7 @@ def needs_offload(target): self.offloaded_states.add(OffloadStateTypeEnum.optim_states) gc.collect() - get_accelerator().empty_cache() + record_empty_cache("stage3.offload-states", get_accelerator().empty_cache) def reload_states(self, non_blocking: bool = False): diff --git a/deepspeed/utils/allocator_telemetry.py b/deepspeed/utils/allocator_telemetry.py new file mode 100644 index 000000000000..c853f236dbc1 --- /dev/null +++ b/deepspeed/utils/allocator_telemetry.py @@ -0,0 +1,183 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Small opt-in allocator diagnostics for the DeepCompile memory experiment.""" + +from collections import Counter +import json +import os +from typing import Any, Callable, Dict, Optional + +MAX_OBSERVATIONS = 64 + +_step: Optional[int] = None +_phase: Optional[str] = None +_empty_cache_counts = Counter() +_empty_cache_observations = [] +_scheduler_decisions = [] +_allocator_retry_samples = [] +_allocator_retry_total = 0 +_dropped_observations = Counter() +_closed = False + + +def enabled() -> bool: + return os.environ.get("DEEPSPEED_ALLOCATOR_TELEMETRY", "0") in {"1", "true", "True"} + + +def memory_stats() -> Dict[str, Any]: + result = { + "num_alloc_retries": None, + "allocated_bytes": None, + "allocated_peak_bytes": None, + "reserved_bytes": None, + "reserved_peak_bytes": None, + "inactive_split_bytes": None, + } + if not enabled(): + return result + try: + from deepspeed.accelerator import get_accelerator + + stats = get_accelerator().memory_stats() or {} + mapping = { + "num_alloc_retries": "num_alloc_retries", + "allocated_bytes": "allocated_bytes.all.current", + "allocated_peak_bytes": "allocated_bytes.all.peak", + "reserved_bytes": "reserved_bytes.all.current", + "reserved_peak_bytes": "reserved_bytes.all.peak", + "inactive_split_bytes": "inactive_split_bytes.all.current", + } + for output_name, stats_name in mapping.items(): + value = stats.get(stats_name) + result[output_name] = int(value) if value is not None else None + except Exception as exc: + result["error"] = type(exc).__name__ + return result + + +def set_step(step: Optional[int], phase: Optional[str] = None) -> None: + global _step, _phase + if enabled(): + _step = step + _phase = phase + + +def _append(kind: str, collection: list, observation: Dict[str, Any]) -> None: + if len(collection) < MAX_OBSERVATIONS: + collection.append(observation) + else: + _dropped_observations[kind] += 1 + + +def record_empty_cache(site: str, callable_: Callable[..., Any], *args, **kwargs): + if not enabled(): + return callable_(*args, **kwargs) + before = memory_stats() + _empty_cache_counts[site] += 1 + error = None + try: + return callable_(*args, **kwargs) + except BaseException as exc: + error = type(exc).__name__ + raise + finally: + _append( + "empty_cache", _empty_cache_observations, { + "site": site, + "ordinal": _empty_cache_counts[site], + "step": _step, + "phase": _phase, + "before": before, + "after": memory_stats(), + "error": error, + }) + + +def record_scheduler_decision(*, + graph_id: int, + backward: bool, + budget_source: Optional[str], + disabled_reason: Optional[str], + max_gathered_bytes: Optional[int], + max_live_gathered_bytes: int, + budget_rejections: int, + over_budget_fallbacks: int, + minimum_rejected_candidate_peak_bytes: Optional[int] = None) -> None: + if not enabled(): + return + _append( + "scheduler", _scheduler_decisions, { + "graph_id": + int(graph_id), + "backward": + bool(backward), + "budget_source": + budget_source, + "disabled_reason": + disabled_reason, + "max_gathered_bytes": + int(max_gathered_bytes) if max_gathered_bytes is not None else None, + "max_live_gathered_bytes": + int(max_live_gathered_bytes), + "budget_rejections": + int(budget_rejections), + "over_budget_fallbacks": + int(over_budget_fallbacks), + "minimum_rejected_candidate_peak_bytes": (int(minimum_rejected_candidate_peak_bytes) + if minimum_rejected_candidate_peak_bytes is not None else None), + }) + + +def record_allocator_retry_sample(tracker_before: int, tracker_after: int) -> None: + global _allocator_retry_total + if not enabled(): + return + delta = max(0, int(tracker_after) - int(tracker_before)) + _allocator_retry_total += delta + if delta: + _append( + "allocator_retry", _allocator_retry_samples, { + "step": _step, + "phase": _phase, + "tracker_before": int(tracker_before), + "tracker_after": int(tracker_after), + "retry_delta": delta, + }) + + +def summary() -> Dict[str, Any]: + return { + "rank": int(os.environ.get("RANK", "0")), + "world_size": int(os.environ.get("WORLD_SIZE", "1")), + "last_step": _step, + "last_phase": _phase, + "final_memory": memory_stats(), + "empty_cache_counts": dict(sorted(_empty_cache_counts.items())), + "empty_cache_observations": list(_empty_cache_observations), + "scheduler_decisions": list(_scheduler_decisions), + "allocator_retry_total": _allocator_retry_total, + "allocator_retry_samples": list(_allocator_retry_samples), + "dropped_observations": dict(sorted(_dropped_observations.items())), + } + + +def close_recorder() -> None: + global _closed + if enabled() and not _closed: + print("DEEPSPEED_ALLOCATOR_TELEMETRY_SUMMARY " + json.dumps(summary(), sort_keys=True), flush=True) + _closed = True + + +def _reset_for_tests() -> None: + global _step, _phase, _allocator_retry_total, _closed + _step = None + _phase = None + _empty_cache_counts.clear() + _empty_cache_observations.clear() + _scheduler_decisions.clear() + _allocator_retry_samples.clear() + _allocator_retry_total = 0 + _dropped_observations.clear() + _closed = False diff --git a/tests/torch_compile/test_deepcompile_z3_release.py b/tests/torch_compile/test_deepcompile_z3_release.py index e9eba252d0a0..d4ce474d9315 100644 --- a/tests/torch_compile/test_deepcompile_z3_release.py +++ b/tests/torch_compile/test_deepcompile_z3_release.py @@ -25,21 +25,25 @@ class TestDeepCompileZ3ReleaseStorage(DistributedTest): def _device(self): return torch.device(get_accelerator().current_device_name()) - def _init_dc(self): + def _init_dc(self, pool_budget=1 << 20): dc = get_deepcompile_handle() dc.init(dist.get_world_group(), CompileConfig(deepcompile=True), 1024) + if pool_budget is not None: + dc.set_z3_gather_buffer_pool_budget_for_test(pool_budget) return dc - def _register_param(self, dc, graph_id, ds_id, shape, persistent=False): + def _register_param(self, dc, graph_id, ds_id, shape, persistent=False, register_graph=True, dtype=torch.float32): device = self._device() world_size = dist.get_world_size() true_numel = math.prod(shape) shard_numel = math.ceil(true_numel / world_size) rank = dist.get_rank() - values = torch.arange(rank * shard_numel, (rank + 1) * shard_numel, device=device, dtype=torch.float32) + values = torch.arange(rank * shard_numel, (rank + 1) * shard_numel, device=device, + dtype=torch.float32).to(dtype) grad_buffer = torch.zeros_like(values) dc.register_z3_param(ds_id, list(shape), values, grad_buffer, persistent, values.dtype) - dc.register_graph_z3(graph_id, [ds_id]) + if register_graph: + dc.register_graph_z3(graph_id, [ds_id]) return values def _gather_view_and_storage(self, shard, graph_id, ds_id): @@ -63,14 +67,28 @@ def _expected_view_sum(self, shape): values = values[:math.prod(shape)].reshape(-1) return values.narrow(0, 0, values.numel() - 1).sum() - def test_storage_resized_to_zero_after_release_single_use(self): - graph_id, ds_id = 9010, 9011 + def _pool_state(self, dc): + keys = ("budget", "charged", "high_water", "entries", "checked_out", "retries", "enabled", "initialized", + "idle_pressure_score", "pressure_recovery_complete", "pressure_recovery_budget", + "pressure_recovery_pending_entries", "pressure_recovery_in_progress") + return dict(zip(keys, dc.get_z3_gather_buffer_pool_state_for_test())) + + def test_storage_reused_after_release_single_use(self): + graph_id, ds_id, next_ds_id = 9010, 9011, 9012 dc = self._init_dc() try: - shard = self._register_param(dc, graph_id, ds_id, [4097]) + shard = self._register_param(dc, graph_id, ds_id, [4097], register_graph=False) + next_shard = self._register_param(dc, graph_id, next_ds_id, [2049], register_graph=False) + dc.register_graph_z3(graph_id, [ds_id, next_ds_id]) view, storage = self._gather_view_and_storage(shard, graph_id, ds_id) + before_ptr = storage.data_ptr() self._release(view, graph_id, ds_id, 1) - assert storage.nbytes() == 0 + assert storage.nbytes() > 0 + + next_view, next_storage = self._gather_view_and_storage(next_shard, graph_id, next_ds_id) + assert next_storage.data_ptr() == before_ptr + assert torch.allclose(next_view.sum(), self._expected_view_sum([2049])) + self._release(next_view, graph_id, next_ds_id, 1) finally: dc.cleanup() @@ -84,10 +102,947 @@ def test_storage_nonzero_until_final_release_when_multi_use(self): self._release(view, graph_id, ds_id, 2) assert storage.nbytes() == before_release_nbytes self._release(view, graph_id, ds_id, 2) + assert storage.nbytes() == before_release_nbytes + finally: + dc.cleanup() + + def test_pool_budget_counts_checked_out_storage(self): + graph_id = 9050 + first_ds_id, checked_out_ds_id, overlapping_ds_id = 9051, 9052, 9053 + dc = self._init_dc(pool_budget=20_000) + try: + first_shard = self._register_param(dc, graph_id, first_ds_id, [4097], register_graph=False) + checked_out_shard = self._register_param(dc, graph_id, checked_out_ds_id, [2049], register_graph=False) + overlapping_shard = self._register_param(dc, graph_id, overlapping_ds_id, [1025], register_graph=False) + dc.register_graph_z3(graph_id, [first_ds_id, checked_out_ds_id, overlapping_ds_id]) + + first_view, first_storage = self._gather_view_and_storage(first_shard, graph_id, first_ds_id) + pool_ptr = first_storage.data_ptr() + self._release(first_view, graph_id, first_ds_id, 1) + + checked_out_view, checked_out_storage = self._gather_view_and_storage(checked_out_shard, graph_id, + checked_out_ds_id) + assert checked_out_storage.data_ptr() == pool_ptr + + overlapping_view, overlapping_storage = self._gather_view_and_storage(overlapping_shard, graph_id, + overlapping_ds_id) + assert overlapping_storage.data_ptr() != pool_ptr + self._release(overlapping_view, graph_id, overlapping_ds_id, 1) + assert overlapping_storage.nbytes() == 0 + + self._release(checked_out_view, graph_id, checked_out_ds_id, 1) + assert checked_out_storage.data_ptr() == pool_ptr + assert checked_out_storage.nbytes() == first_storage.nbytes() + finally: + dc.cleanup() + + def test_zero_pool_budget_uses_resize_to_zero_fallback(self): + graph_id, ds_id = 9060, 9061 + dc = self._init_dc(pool_budget=0) + try: + shard = self._register_param(dc, graph_id, ds_id, [4097]) + view, storage = self._gather_view_and_storage(shard, graph_id, ds_id) + self._release(view, graph_id, ds_id, 1) assert storage.nbytes() == 0 finally: dc.cleanup() + def test_prefetched_storage_is_not_admitted_to_demand_gather_pool(self): + graph_id, prefetched_ds_id, demand_ds_id = 9070, 9071, 9072 + dc = self._init_dc() + try: + prefetched_shard = self._register_param(dc, graph_id, prefetched_ds_id, [4097], register_graph=False) + demand_shard = self._register_param(dc, graph_id, demand_ds_id, [2049], register_graph=False) + dc.register_graph_z3(graph_id, [prefetched_ds_id, demand_ds_id]) + + torch.ops.dc.prefetch_params_fused.default(graph_id, [prefetched_shard], [prefetched_ds_id]) + prefetched_view, prefetched_storage = self._gather_view_and_storage(prefetched_shard, graph_id, + prefetched_ds_id) + self._release(prefetched_view, graph_id, prefetched_ds_id, 1) + assert prefetched_storage.nbytes() == 0 + + demand_view, demand_storage = self._gather_view_and_storage(demand_shard, graph_id, demand_ds_id) + self._release(demand_view, graph_id, demand_ds_id, 1) + assert demand_storage.nbytes() > 0 + finally: + dc.cleanup() + + def test_prefetch_excludes_existing_pool_storage_from_demand_reuse(self): + graph_id = 9090 + first_ds_id, prefetched_ds_id, demand_ds_id = 9091, 9092, 9093 + dc = self._init_dc() + try: + first_shard = self._register_param(dc, graph_id, first_ds_id, [4097], register_graph=False) + prefetched_shard = self._register_param(dc, graph_id, prefetched_ds_id, [2049], register_graph=False) + demand_shard = self._register_param(dc, graph_id, demand_ds_id, [1025], register_graph=False) + dc.register_graph_z3(graph_id, [first_ds_id, prefetched_ds_id, demand_ds_id]) + + first_view, first_storage = self._gather_view_and_storage(first_shard, graph_id, first_ds_id) + pool_ptr = first_storage.data_ptr() + self._release(first_view, graph_id, first_ds_id, 1) + + prefetched_view, prefetched_storage = self._gather_view_and_storage(prefetched_shard, graph_id, + prefetched_ds_id) + assert prefetched_storage.data_ptr() == pool_ptr + dc.set_z3_param_valid_for_test(prefetched_ds_id, False) + torch.ops.dc.prefetch_params_fused.default(graph_id, [prefetched_shard], [prefetched_ds_id]) + + self._release(prefetched_view, graph_id, prefetched_ds_id, 1) + assert prefetched_storage.nbytes() == 0 + + demand_view, demand_storage = self._gather_view_and_storage(demand_shard, graph_id, demand_ds_id) + self._release(demand_view, graph_id, demand_ds_id, 1) + assert demand_storage.nbytes() > 0 + finally: + dc.cleanup() + + def test_prefetch_preparation_failure_rolls_back_storage_exclusion(self): + graph_id = 9094 + first_ds_id, prefetched_ds_id, demand_ds_id = 9095, 9096, 9097 + dc = self._init_dc() + try: + first_shard = self._register_param(dc, graph_id, first_ds_id, [4097], register_graph=False) + prefetched_shard = self._register_param(dc, graph_id, prefetched_ds_id, [2049], register_graph=False) + demand_shard = self._register_param(dc, graph_id, demand_ds_id, [1025], register_graph=False) + dc.register_graph_z3(graph_id, [first_ds_id, prefetched_ds_id, demand_ds_id]) + + first_view, first_storage = self._gather_view_and_storage(first_shard, graph_id, first_ds_id) + pool_ptr = first_storage.data_ptr() + self._release(first_view, graph_id, first_ds_id, 1) + + prefetched_view, prefetched_storage = self._gather_view_and_storage(prefetched_shard, graph_id, + prefetched_ds_id) + assert prefetched_storage.data_ptr() == pool_ptr + dc.set_z3_param_valid_for_test(prefetched_ds_id, False) + dc.set_z3_prefetch_fail_after_exclusions_for_test(1) + with pytest.raises(RuntimeError, match="injected prefetch preparation failure"): + torch.ops.dc.prefetch_params_fused.default(graph_id, [prefetched_shard], [prefetched_ds_id]) + dc.set_z3_prefetch_fail_after_exclusions_for_test(0) + + self._release(prefetched_view, graph_id, prefetched_ds_id, 1) + assert prefetched_storage.nbytes() > 0 + demand_view, demand_storage = self._gather_view_and_storage(demand_shard, graph_id, demand_ds_id) + assert demand_storage.data_ptr() == pool_ptr + self._release(demand_view, graph_id, demand_ds_id, 1) + finally: + dc.cleanup() + + def test_pressure_recovery_retires_checked_out_working_set_on_final_returns(self): + graph_id = 9100 + large_ds_id, small_ds_id, transient_ds_id = 9101, 9102, 9103 + dc = self._init_dc(pool_budget=None) + try: + large_shard = self._register_param(dc, graph_id, large_ds_id, [1_048_577], register_graph=False) + small_shard = self._register_param(dc, graph_id, small_ds_id, [524_289], register_graph=False) + transient_shard = self._register_param(dc, graph_id, transient_ds_id, [262_145], register_graph=False) + dc.register_graph_z3(graph_id, [large_ds_id, small_ds_id, transient_ds_id]) + + gib = 1 << 30 + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(0, gib, 8 * gib) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(1, gib, 8 * gib) + + # Prime two distinct pool entries, then check both out together so + # recovery cannot depend on an all-idle observation. + large_view, large_storage = self._gather_view_and_storage(large_shard, graph_id, large_ds_id) + small_view, small_storage = self._gather_view_and_storage(small_shard, graph_id, small_ds_id) + self._release(large_view, graph_id, large_ds_id, 1) + self._release(small_view, graph_id, small_ds_id, 1) + large_capacity = large_storage.nbytes() + small_capacity = small_storage.nbytes() + assert large_capacity > small_capacity + + checked_out_large, checked_out_large_storage = self._gather_view_and_storage( + large_shard, graph_id, large_ds_id) + checked_out_small, checked_out_small_storage = self._gather_view_and_storage( + small_shard, graph_id, small_ds_id) + + before_pressure = self._pool_state(dc) + assert before_pressure["entries"] == 2 + assert before_pressure["checked_out"] == 2 + assert before_pressure["charged"] == large_capacity + small_capacity + + # Isolated retries preserve the byte-exact working set. Crossing + # the sustained-pressure threshold latches recovery without + # reclaiming either live lease. + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(2, 8 << 20, 8 * gib) + after_pressure = self._pool_state(dc) + assert after_pressure["budget"] == before_pressure["charged"] + assert after_pressure["charged"] == before_pressure["charged"] + assert after_pressure["idle_pressure_score"] == 2 + + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(3, 8 << 20, 8 * gib) + after_threshold = self._pool_state(dc) + assert after_threshold["entries"] == 2 + assert after_threshold["checked_out"] == 2 + assert after_threshold["charged"] == before_pressure["charged"] + assert after_threshold["idle_pressure_score"] == 0 + assert after_threshold["pressure_recovery_in_progress"] == 1 + assert after_threshold["pressure_recovery_complete"] == 0 + assert after_threshold["pressure_recovery_budget"] == before_pressure["charged"] + + # New demand during the drain cannot extend pool ownership or the + # captured recovery multiset. + transient_view, transient_storage = self._gather_view_and_storage(transient_shard, graph_id, + transient_ds_id) + self._release(transient_view, graph_id, transient_ds_id, 1) + during_drain = self._pool_state(dc) + assert during_drain["entries"] == 2 + assert during_drain["checked_out"] == 2 + assert during_drain["pressure_recovery_budget"] == before_pressure["charged"] + assert transient_storage.nbytes() == 0 + + # Each final return makes bounded progress. No new pressure sample + # or all-idle instant is needed to finish the one-shot drain. + self._release(checked_out_small, graph_id, small_ds_id, 1) + partial = self._pool_state(dc) + assert partial["entries"] == 1 + assert partial["checked_out"] == 1 + assert partial["charged"] == large_capacity + assert partial["pressure_recovery_in_progress"] == 1 + assert partial["pressure_recovery_complete"] == 0 + assert partial["pressure_recovery_pending_entries"] == 1 + assert checked_out_small_storage.nbytes() == 0 + + self._release(checked_out_large, graph_id, large_ds_id, 1) + drained = self._pool_state(dc) + assert drained["entries"] == 0 + assert drained["charged"] == 0 + assert drained["budget"] == before_pressure["charged"] + assert drained["enabled"] == 1 + assert drained["pressure_recovery_in_progress"] == 0 + assert drained["pressure_recovery_complete"] == 1 + assert drained["pressure_recovery_pending_entries"] == 2 + assert checked_out_large_storage.nbytes() == 0 + + # The complete typed/capacity/device multiset is admitted again. + recovered_small, recovered_small_storage = self._gather_view_and_storage( + small_shard, graph_id, small_ds_id) + self._release(recovered_small, graph_id, small_ds_id, 1) + recovered_large, recovered_large_storage = self._gather_view_and_storage( + large_shard, graph_id, large_ds_id) + self._release(recovered_large, graph_id, large_ds_id, 1) + recovered_state = self._pool_state(dc) + assert recovered_state["entries"] == 2 + assert recovered_state["charged"] == before_pressure["charged"] + assert recovered_state["pressure_recovery_pending_entries"] == 0 + assert recovered_small_storage.nbytes() == small_capacity + assert recovered_large_storage.nbytes() == large_capacity + finally: + dc.cleanup() + + def test_pressure_recovery_retires_multi_user_lease_only_after_final_return(self): + graph_id, ds_id = 9170, 9171 + dc = self._init_dc(pool_budget=None) + try: + shard = self._register_param(dc, graph_id, ds_id, [1_048_577]) + gib = 1 << 30 + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(0, gib, 8 * gib) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(1, gib, 8 * gib) + + first_view, _ = self._gather_view_and_storage(shard, graph_id, ds_id) + self._release(first_view, graph_id, ds_id, 1) + held_view, held_storage = self._gather_view_and_storage(shard, graph_id, ds_id) + held_capacity = held_storage.nbytes() + + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(5, 8 << 20, 8 * gib) + assert self._pool_state(dc)["pressure_recovery_in_progress"] == 1 + + self._release(held_view, graph_id, ds_id, 2) + first_return = self._pool_state(dc) + assert first_return["entries"] == 1 + assert first_return["checked_out"] == 1 + assert first_return["charged"] == held_capacity + assert first_return["pressure_recovery_in_progress"] == 1 + assert held_storage.nbytes() == held_capacity + + self._release(held_view, graph_id, ds_id, 2) + final_return = self._pool_state(dc) + assert final_return["entries"] == 0 + assert final_return["charged"] == 0 + assert final_return["pressure_recovery_in_progress"] == 0 + assert final_return["pressure_recovery_complete"] == 1 + assert held_storage.nbytes() == 0 + finally: + dc.cleanup() + + def test_selective_persistence_detaches_reused_storage_from_pool_accounting(self): + graph_id, second_graph_id, pooled_ds_id, persistent_ds_id = 91006, 91007, 91008, 91009 + dc = self._init_dc(pool_budget=None) + try: + pooled_shard = self._register_param(dc, graph_id, pooled_ds_id, [4097], register_graph=False) + persistent_shard = self._register_param(dc, graph_id, persistent_ds_id, [2049], register_graph=False) + dc.register_graph_z3(graph_id, [pooled_ds_id, persistent_ds_id]) + dc.register_graph_z3(second_graph_id, [persistent_ds_id]) + gib = 1 << 30 + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(0, gib, 8 * gib) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(1, gib, 8 * gib) + + pooled_view, pooled_storage = self._gather_view_and_storage(pooled_shard, graph_id, pooled_ds_id) + pooled_ptr = pooled_storage.data_ptr() + self._release(pooled_view, graph_id, pooled_ds_id, 1) + assert self._pool_state(dc)["entries"] == 1 + + dc.set_persistent(persistent_ds_id) + persistent_view, persistent_storage = self._gather_view_and_storage(persistent_shard, graph_id, + persistent_ds_id) + detached = self._pool_state(dc) + assert persistent_storage.data_ptr() == pooled_ptr + assert persistent_storage.nbytes() > 0 + assert detached["entries"] == 0 + assert detached["checked_out"] == 0 + assert detached["charged"] == 0 + + dc.set_persistent(persistent_ds_id) + assert persistent_storage.nbytes() > 0 + assert self._pool_state(dc)["charged"] == 0 + + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(5, 8 << 20, 8 * gib) + # Keep a tensor dependency so the dispatcher can materialize the + # Any argument; an empty Python list has no inferable element type. + torch.ops.dc.end_backward.default([torch.empty(0, device=persistent_view.device)], graph_id, True) + assert persistent_storage.nbytes() > 0 + assert torch.allclose(persistent_view.sum(), self._expected_view_sum([2049])) + finally: + dc.cleanup() + + def test_pressure_recovery_exclusion_retires_without_phantom_target(self): + graph_id, ds_id = 9172, 9173 + dc = self._init_dc(pool_budget=None) + try: + shard = self._register_param(dc, graph_id, ds_id, [1_048_577]) + gib = 1 << 30 + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(0, gib, 8 * gib) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(1, gib, 8 * gib) + + first_view, _ = self._gather_view_and_storage(shard, graph_id, ds_id) + self._release(first_view, graph_id, ds_id, 1) + held_view, held_storage = self._gather_view_and_storage(shard, graph_id, ds_id) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(5, 8 << 20, 8 * gib) + + dc.set_z3_param_valid_for_test(ds_id, False) + torch.ops.dc.prefetch_params_fused.default(graph_id, [shard], [ds_id]) + self._release(held_view, graph_id, ds_id, 1) + excluded = self._pool_state(dc) + assert excluded["entries"] == 0 + assert excluded["charged"] == 0 + assert excluded["pressure_recovery_in_progress"] == 0 + assert excluded["pressure_recovery_complete"] == 1 + assert excluded["pressure_recovery_budget"] == 0 + assert excluded["pressure_recovery_pending_entries"] == 0 + assert held_storage.nbytes() == 0 + finally: + dc.cleanup() + + def test_pressure_recovery_discard_completes_without_phantom_target(self): + graph_id, ds_id = 9174, 9175 + dc = self._init_dc(pool_budget=None) + try: + shard = self._register_param(dc, graph_id, ds_id, [1_048_577]) + gib = 1 << 30 + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(0, gib, 8 * gib) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(1, gib, 8 * gib) + + first_view, _ = self._gather_view_and_storage(shard, graph_id, ds_id) + self._release(first_view, graph_id, ds_id, 1) + held_view, held_storage = self._gather_view_and_storage(shard, graph_id, ds_id) + held_capacity = held_storage.nbytes() + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(5, 8 << 20, 8 * gib) + + dc.invalidate_gathered_param(ds_id) + discarded = self._pool_state(dc) + assert discarded["entries"] == 0 + assert discarded["charged"] == 0 + assert discarded["pressure_recovery_in_progress"] == 0 + assert discarded["pressure_recovery_complete"] == 1 + assert discarded["pressure_recovery_budget"] == 0 + assert discarded["pressure_recovery_pending_entries"] == 0 + # Discard drops pool ownership but must not mutate a still-live + # profiling alias. + assert held_storage.nbytes() == held_capacity + del held_view + finally: + dc.cleanup() + + def test_repeated_allocator_pressure_recovers_once_and_readmits_non_aligned_working_set(self): + graph_id, ds_id = 9103, 9104 + dc = self._init_dc(pool_budget=None) + try: + shard = self._register_param(dc, graph_id, ds_id, [1_048_577]) + gib = 1 << 30 + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(0, gib, 8 * gib) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(1, gib, 8 * gib) + + view, storage = self._gather_view_and_storage(shard, graph_id, ds_id) + self._release(view, graph_id, ds_id, 1) + assert storage.nbytes() > 2 << 20 + + before_pressure = self._pool_state(dc) + assert before_pressure["charged"] % (2 << 20) != 0 + + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(2, 8 << 20, 8 * gib) + after_pressure = self._pool_state(dc) + assert after_pressure["budget"] == before_pressure["charged"] + assert after_pressure["charged"] == before_pressure["charged"] + assert after_pressure["idle_pressure_score"] == 2 + assert storage.nbytes() == before_pressure["charged"] + + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(3, 8 << 20, 8 * gib) + after_threshold = self._pool_state(dc) + assert after_threshold["budget"] == before_pressure["charged"] + assert after_threshold["charged"] == 0 + assert after_threshold["entries"] == 0 + assert after_threshold["enabled"] == 1 + assert after_threshold["idle_pressure_score"] == 0 + assert after_threshold["pressure_recovery_complete"] == 1 + assert storage.nbytes() == 0 + + recovered_view, recovered_storage = self._gather_view_and_storage(shard, graph_id, ds_id) + self._release(recovered_view, graph_id, ds_id, 1) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(100, 8 << 20, 8 * gib) + after_repeat = self._pool_state(dc) + assert after_repeat["entries"] == 1 + assert after_repeat["charged"] == before_pressure["charged"] + assert after_repeat["budget"] == before_pressure["charged"] + assert after_repeat["pressure_recovery_complete"] == 1 + assert recovered_storage.nbytes() == before_pressure["charged"] + finally: + dc.cleanup() + + def test_repeated_allocator_pressure_preserves_idle_pool_below_budget(self): + graph_id, ds_id = 9108, 9109 + dc = self._init_dc(pool_budget=None) + try: + shard = self._register_param(dc, graph_id, ds_id, [1_048_577]) + gib = 1 << 30 + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(0, gib, 8 * gib) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(1, gib, 8 * gib) + + view, storage = self._gather_view_and_storage(shard, graph_id, ds_id) + self._release(view, graph_id, ds_id, 1) + before_pressure = self._pool_state(dc) + assert before_pressure["charged"] < before_pressure["budget"] + + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(2, gib, 8 * gib) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(3, gib, 8 * gib) + below_budget = self._pool_state(dc) + assert below_budget["charged"] == before_pressure["charged"] + assert below_budget["budget"] == before_pressure["budget"] + assert below_budget["entries"] == 1 + assert below_budget["enabled"] == 1 + assert below_budget["idle_pressure_score"] == 3 + assert storage.nbytes() == before_pressure["charged"] + + # The same sustained score performs one recovery once pressure + # lowers the budget to the retained charge. + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(4, 8 << 20, 8 * gib) + at_budget = self._pool_state(dc) + assert at_budget["charged"] == 0 + assert at_budget["budget"] == before_pressure["charged"] + assert at_budget["entries"] == 0 + assert at_budget["enabled"] == 1 + assert at_budget["idle_pressure_score"] == 0 + assert at_budget["pressure_recovery_complete"] == 1 + assert storage.nbytes() == 0 + + recovered_view, recovered_storage = self._gather_view_and_storage(shard, graph_id, ds_id) + self._release(recovered_view, graph_id, ds_id, 1) + assert self._pool_state(dc)["charged"] == before_pressure["charged"] + assert recovered_storage.nbytes() == before_pressure["charged"] + finally: + dc.cleanup() + + def test_allocator_retry_jump_recovers_once_and_preserves_complete_hot_working_set(self): + graph_id, large_ds_id, small_ds_id = 9105, 9106, 9107 + dc = self._init_dc(pool_budget=None) + try: + large_shard = self._register_param(dc, graph_id, large_ds_id, [1_048_577], register_graph=False) + small_shard = self._register_param(dc, graph_id, small_ds_id, [524_289], register_graph=False) + dc.register_graph_z3(graph_id, [large_ds_id, small_ds_id]) + gib = 1 << 30 + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(0, gib, 8 * gib) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(1, gib, 8 * gib) + + large_view, large_storage = self._gather_view_and_storage(large_shard, graph_id, large_ds_id) + small_view, small_storage = self._gather_view_and_storage(small_shard, graph_id, small_ds_id) + self._release(large_view, graph_id, large_ds_id, 1) + self._release(small_view, graph_id, small_ds_id, 1) + before_jump = self._pool_state(dc) + assert before_jump["entries"] == 2 + assert large_storage.nbytes() > small_storage.nbytes() + large_capacity = large_storage.nbytes() + small_capacity = small_storage.nbytes() + working_set_capacity = before_jump["charged"] + + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(5, 8 << 20, 8 * gib) + after_jump = self._pool_state(dc) + assert after_jump["entries"] == 0 + assert after_jump["charged"] == 0 + assert after_jump["budget"] == working_set_capacity + assert after_jump["enabled"] == 1 + assert after_jump["idle_pressure_score"] == 0 + assert after_jump["pressure_recovery_complete"] == 1 + assert after_jump["pressure_recovery_budget"] == working_set_capacity + assert after_jump["pressure_recovery_pending_entries"] == 2 + assert large_storage.nbytes() == 0 + assert small_storage.nbytes() == 0 + + # A retry in the empty recovery window cannot shrink the budget + # below the remembered hot-working-set floor. + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(6, 8 << 20, 8 * gib) + empty_window = self._pool_state(dc) + assert empty_window["entries"] == 0 + assert empty_window["budget"] == working_set_capacity + assert empty_window["pressure_recovery_budget"] == working_set_capacity + assert empty_window["pressure_recovery_pending_entries"] == 2 + + # Both exact typed/device storage identities are admitted. The + # smaller entry can arrive first without consuming the larger + # target's identity or blocking its later admission. + recovered_small_view, recovered_small_storage = self._gather_view_and_storage( + small_shard, graph_id, small_ds_id) + self._release(recovered_small_view, graph_id, small_ds_id, 1) + after_small_readmit = self._pool_state(dc) + assert after_small_readmit["entries"] == 1 + assert after_small_readmit["charged"] == small_capacity + assert after_small_readmit["pressure_recovery_pending_entries"] == 1 + + recovered_view, recovered_storage = self._gather_view_and_storage(large_shard, graph_id, large_ds_id) + recovered_ptr = recovered_storage.data_ptr() + self._release(recovered_view, graph_id, large_ds_id, 1) + after_readmit = self._pool_state(dc) + assert after_readmit["entries"] == 2 + assert after_readmit["charged"] == working_set_capacity + assert after_readmit["pressure_recovery_pending_entries"] == 0 + assert recovered_storage.nbytes() == large_capacity + + # Counter regression clears the score but not the lifecycle latch + # or floor; later retry waves cannot evict the re-admitted working set. + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(0, 8 << 20, 8 * gib) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(8, 8 << 20, 8 * gib) + after_repeat = self._pool_state(dc) + assert after_repeat["entries"] == 2 + assert after_repeat["charged"] == working_set_capacity + assert after_repeat["budget"] == working_set_capacity + assert after_repeat["enabled"] == 1 + assert after_repeat["idle_pressure_score"] == 8 + assert after_repeat["pressure_recovery_complete"] == 1 + assert after_repeat["pressure_recovery_budget"] == working_set_capacity + assert after_repeat["pressure_recovery_pending_entries"] == 0 + assert recovered_storage.nbytes() == large_capacity + + reused_view, reused_storage = self._gather_view_and_storage(large_shard, graph_id, large_ds_id) + assert reused_storage.data_ptr() == recovered_ptr + self._release(reused_view, graph_id, large_ds_id, 1) + + dc.cleanup() + dc = self._init_dc(pool_budget=None) + self._register_param(dc, 9108, 9110, [3]) + reset_lifecycle = self._pool_state(dc) + assert reset_lifecycle["pressure_recovery_complete"] == 0 + assert reset_lifecycle["pressure_recovery_budget"] == 0 + assert reset_lifecycle["pressure_recovery_pending_entries"] == 0 + finally: + dc.cleanup() + + def test_recovery_multiset_admits_duplicate_targets_once_each(self): + graph_id = 9113 + first_ds_id, second_ds_id, extra_ds_id = 9114, 9115, 9116 + dc = self._init_dc(pool_budget=None) + try: + first_shard = self._register_param(dc, graph_id, first_ds_id, [524_289], register_graph=False) + second_shard = self._register_param(dc, graph_id, second_ds_id, [524_289], register_graph=False) + extra_shard = self._register_param(dc, graph_id, extra_ds_id, [524_289], register_graph=False) + dc.register_graph_z3(graph_id, [first_ds_id, second_ds_id, extra_ds_id]) + gib = 1 << 30 + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(0, gib, 8 * gib) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(1, gib, 8 * gib) + + first_view, first_storage = self._gather_view_and_storage(first_shard, graph_id, first_ds_id) + second_view, second_storage = self._gather_view_and_storage(second_shard, graph_id, second_ds_id) + assert first_storage.data_ptr() != second_storage.data_ptr() + self._release(first_view, graph_id, first_ds_id, 1) + self._release(second_view, graph_id, second_ds_id, 1) + before_recovery = self._pool_state(dc) + assert before_recovery["entries"] == 2 + assert first_storage.nbytes() == second_storage.nbytes() + + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(5, 8 << 20, 8 * gib) + recovery = self._pool_state(dc) + assert recovery["entries"] == 0 + assert recovery["pressure_recovery_budget"] == before_recovery["charged"] + assert recovery["pressure_recovery_pending_entries"] == 2 + + recovered_first, _ = self._gather_view_and_storage(first_shard, graph_id, first_ds_id) + recovered_second, _ = self._gather_view_and_storage(second_shard, graph_id, second_ds_id) + self._release(recovered_first, graph_id, first_ds_id, 1) + self._release(recovered_second, graph_id, second_ds_id, 1) + recovered = self._pool_state(dc) + assert recovered["entries"] == 2 + assert recovered["charged"] == before_recovery["charged"] + assert recovered["pressure_recovery_pending_entries"] == 0 + + held_first, _ = self._gather_view_and_storage(first_shard, graph_id, first_ds_id) + held_second, _ = self._gather_view_and_storage(second_shard, graph_id, second_ds_id) + extra_view, extra_storage = self._gather_view_and_storage(extra_shard, graph_id, extra_ds_id) + self._release(extra_view, graph_id, extra_ds_id, 1) + assert extra_storage.nbytes() == 0 + assert self._pool_state(dc)["entries"] == 2 + self._release(held_first, graph_id, first_ds_id, 1) + self._release(held_second, graph_id, second_ds_id, 1) + assert self._pool_state(dc)["pressure_recovery_pending_entries"] == 0 + finally: + dc.cleanup() + + def test_recovery_identity_rejects_equal_byte_incompatible_dtype(self): + graph_id, hot_ds_id, incompatible_ds_id = 9120, 9121, 9122 + dc = self._init_dc(pool_budget=None) + try: + hot_shard = self._register_param(dc, + graph_id, + hot_ds_id, [1_048_576], + register_graph=False, + dtype=torch.float32) + incompatible_shard = self._register_param(dc, + graph_id, + incompatible_ds_id, [2_097_152], + register_graph=False, + dtype=torch.bfloat16) + dc.register_graph_z3(graph_id, [hot_ds_id, incompatible_ds_id]) + gib = 1 << 30 + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(0, gib, 8 * gib) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(1, gib, 8 * gib) + + hot_view, hot_storage = self._gather_view_and_storage(hot_shard, graph_id, hot_ds_id) + self._release(hot_view, graph_id, hot_ds_id, 1) + hot_capacity = hot_storage.nbytes() + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(5, 8 << 20, 8 * gib) + assert self._pool_state(dc)["pressure_recovery_budget"] == hot_capacity + + incompatible_view, incompatible_storage = self._gather_view_and_storage( + incompatible_shard, graph_id, incompatible_ds_id) + assert incompatible_storage.nbytes() == hot_capacity + self._release(incompatible_view, graph_id, incompatible_ds_id, 1) + assert incompatible_storage.nbytes() == 0 + assert self._pool_state(dc)["entries"] == 0 + + held_incompatible_view, held_incompatible_storage = self._gather_view_and_storage( + incompatible_shard, graph_id, incompatible_ds_id) + recovered_hot_view, recovered_hot_storage = self._gather_view_and_storage(hot_shard, graph_id, hot_ds_id) + recovered_hot_ptr = recovered_hot_storage.data_ptr() + self._release(recovered_hot_view, graph_id, hot_ds_id, 1) + recovered = self._pool_state(dc) + assert recovered["entries"] == 1 + assert recovered["charged"] == hot_capacity + + self._release(held_incompatible_view, graph_id, incompatible_ds_id, 1) + assert held_incompatible_storage.nbytes() == 0 + assert self._pool_state(dc)["charged"] == hot_capacity + + reused_hot_view, reused_hot_storage = self._gather_view_and_storage(hot_shard, graph_id, hot_ds_id) + assert reused_hot_storage.data_ptr() == recovered_hot_ptr + self._release(reused_hot_view, graph_id, hot_ds_id, 1) + finally: + dc.cleanup() + + def test_adaptive_hard_cap_discards_checked_out_storage_on_return(self): + graph_id, first_ds_id, checked_out_ds_id = 9130, 9131, 9132 + dc = self._init_dc(pool_budget=None) + try: + first_shard = self._register_param(dc, graph_id, first_ds_id, [1_048_577], register_graph=False) + checked_out_shard = self._register_param(dc, graph_id, checked_out_ds_id, [524_289], register_graph=False) + dc.register_graph_z3(graph_id, [first_ds_id, checked_out_ds_id]) + gib = 1 << 30 + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(0, gib, 8 * gib) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(1, gib, 8 * gib) + + first_view, _ = self._gather_view_and_storage(first_shard, graph_id, first_ds_id) + self._release(first_view, graph_id, first_ds_id, 1) + checked_out_view, checked_out_storage = self._gather_view_and_storage(checked_out_shard, graph_id, + checked_out_ds_id) + + # total / 32 is a 2 MiB hard cap, below this checked-out lease. + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(2, 8 << 20, 64 << 20) + state = self._pool_state(dc) + assert state["budget"] == 2 << 20 + assert state["charged"] > state["budget"] + assert checked_out_storage.nbytes() > state["budget"] + self._release(checked_out_view, graph_id, checked_out_ds_id, 1) + assert checked_out_storage.nbytes() == 0 + finally: + dc.cleanup() + + def test_over_cap_recovery_drains_and_selects_largest_fitting_target(self): + graph_id, large_ds_id, small_ds_id = 9176, 9177, 9178 + dc = self._init_dc(pool_budget=None) + try: + large_shard = self._register_param(dc, graph_id, large_ds_id, [1_048_577], register_graph=False) + small_shard = self._register_param(dc, graph_id, small_ds_id, [524_289], register_graph=False) + dc.register_graph_z3(graph_id, [large_ds_id, small_ds_id]) + gib = 1 << 30 + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(0, gib, 8 * gib) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(1, gib, 8 * gib) + + large_view, large_storage = self._gather_view_and_storage(large_shard, graph_id, large_ds_id) + small_view, small_storage = self._gather_view_and_storage(small_shard, graph_id, small_ds_id) + self._release(large_view, graph_id, large_ds_id, 1) + self._release(small_view, graph_id, small_ds_id, 1) + large_capacity = large_storage.nbytes() + small_capacity = small_storage.nbytes() + assert large_capacity > small_capacity + + # total / 32 is a 6 MiB cap. Both non-aligned targets exceed it + # together, while the larger target fits by itself. + hard_cap = 6 << 20 + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(5, 8 << 20, 192 << 20) + recovered = self._pool_state(dc) + assert large_capacity <= hard_cap + assert large_capacity + small_capacity > hard_cap + assert recovered["entries"] == 0 + assert recovered["charged"] == 0 + assert recovered["pressure_recovery_complete"] == 1 + assert recovered["pressure_recovery_budget"] == large_capacity + assert recovered["pressure_recovery_budget"] <= hard_cap + assert recovered["pressure_recovery_pending_entries"] == 1 + + rejected_small, rejected_small_storage = self._gather_view_and_storage(small_shard, graph_id, small_ds_id) + self._release(rejected_small, graph_id, small_ds_id, 1) + assert rejected_small_storage.nbytes() == 0 + assert self._pool_state(dc)["entries"] == 0 + + admitted_large, admitted_large_storage = self._gather_view_and_storage(large_shard, graph_id, large_ds_id) + self._release(admitted_large, graph_id, large_ds_id, 1) + admitted = self._pool_state(dc) + assert admitted["entries"] == 1 + assert admitted["charged"] == large_capacity + assert admitted["pressure_recovery_pending_entries"] == 0 + assert admitted_large_storage.nbytes() == large_capacity + finally: + dc.cleanup() + + def test_recovery_floor_yields_to_hard_cap_before_hot_readmission(self): + graph_id, ds_id = 9133, 9134 + dc = self._init_dc(pool_budget=None) + try: + shard = self._register_param(dc, graph_id, ds_id, [1_048_577]) + gib = 1 << 30 + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(0, gib, 8 * gib) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(1, gib, 8 * gib) + + view, storage = self._gather_view_and_storage(shard, graph_id, ds_id) + self._release(view, graph_id, ds_id, 1) + hot_capacity = storage.nbytes() + assert hot_capacity > 2 << 20 + + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(5, 8 << 20, 8 * gib) + recovered = self._pool_state(dc) + assert recovered["pressure_recovery_complete"] == 1 + assert recovered["pressure_recovery_budget"] == hot_capacity + assert recovered["entries"] == 0 + + # total / 32 is now 2 MiB. A remembered target larger than the new + # cap is removed so the target multiset remains satisfiable. + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(6, 8 << 20, 64 << 20) + capped = self._pool_state(dc) + assert capped["budget"] == 2 << 20 + assert capped["pressure_recovery_budget"] == 0 + assert capped["pressure_recovery_pending_entries"] == 0 + + hot_view, hot_storage = self._gather_view_and_storage(shard, graph_id, ds_id) + self._release(hot_view, graph_id, ds_id, 1) + after_release = self._pool_state(dc) + assert after_release["entries"] == 0 + assert after_release["charged"] == 0 + assert hot_storage.nbytes() == 0 + finally: + dc.cleanup() + + def test_hard_cap_reclaim_preempts_recovery_below_hot_buffer_size(self): + graph_id, ds_id = 9135, 9136 + dc = self._init_dc(pool_budget=None) + try: + shard = self._register_param(dc, graph_id, ds_id, [1_048_577]) + gib = 1 << 30 + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(0, gib, 8 * gib) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(1, gib, 8 * gib) + + view, storage = self._gather_view_and_storage(shard, graph_id, ds_id) + self._release(view, graph_id, ds_id, 1) + hot_capacity = storage.nbytes() + assert hot_capacity > 2 << 20 + + # The 2 MiB hard cap cannot hold this target. Recovery must still + # drain the entry and close its one-shot lifecycle with no target, + # rather than waiting forever for a budget-fitting all-idle state. + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(5, 8 << 20, 64 << 20) + capped = self._pool_state(dc) + assert capped["budget"] == 0 + assert capped["charged"] == 0 + assert capped["entries"] == 0 + assert capped["idle_pressure_score"] == 0 + assert capped["pressure_recovery_complete"] == 1 + assert capped["pressure_recovery_budget"] == 0 + assert capped["pressure_recovery_pending_entries"] == 0 + assert storage.nbytes() == 0 + + hot_view, hot_storage = self._gather_view_and_storage(shard, graph_id, ds_id) + self._release(hot_view, graph_id, ds_id, 1) + after_release = self._pool_state(dc) + assert after_release["entries"] == 0 + assert after_release["charged"] == 0 + assert after_release["pressure_recovery_complete"] == 1 + assert hot_storage.nbytes() == 0 + finally: + dc.cleanup() + + def test_new_demand_evicts_oldest_idle_entry_after_pressure(self): + graph_id = 9140 + first_ds_id, second_ds_id, demand_ds_id = 9141, 9142, 9143 + dc = self._init_dc(pool_budget=None) + try: + first_shard = self._register_param(dc, graph_id, first_ds_id, [524_289], register_graph=False) + second_shard = self._register_param(dc, graph_id, second_ds_id, [262_145], register_graph=False) + demand_shard = self._register_param(dc, graph_id, demand_ds_id, [786_433], register_graph=False) + dc.register_graph_z3(graph_id, [first_ds_id, second_ds_id, demand_ds_id]) + gib = 1 << 30 + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(0, gib, 8 * gib) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(1, gib, 8 * gib) + + first_view, first_storage = self._gather_view_and_storage(first_shard, graph_id, first_ds_id) + second_view, second_storage = self._gather_view_and_storage(second_shard, graph_id, second_ds_id) + self._release(first_view, graph_id, first_ds_id, 1) + self._release(second_view, graph_id, second_ds_id, 1) + + # Preserve 6 MiB of budget after pressure, then let admission-time + # LRU reclaim the older first entry for larger new demand. + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(2, 24 << 20, 8 * gib) + demand_view, _ = self._gather_view_and_storage(demand_shard, graph_id, demand_ds_id) + self._release(demand_view, graph_id, demand_ds_id, 1) + assert first_storage.nbytes() == 0 + assert second_storage.nbytes() > 0 + finally: + dc.cleanup() + + def test_allocator_retry_reset_and_jump_keep_budget_state_consistent(self): + dc = self._init_dc(pool_budget=None) + try: + # Registering a graph keeps the process-global weak pool alive + # between the individual pressure-seam calls below. + self._register_param(dc, 9150, 9151, [3]) + gib = 1 << 30 + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(10, gib, 8 * gib) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(11, gib, 8 * gib) + initial_state = self._pool_state(dc) + assert initial_state["budget"] == 256 << 20 + assert initial_state["idle_pressure_score"] == 1 + + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(0, 8 << 20, 8 * gib) + reset_state = self._pool_state(dc) + assert reset_state["retries"] == 0 + assert reset_state["budget"] == 256 << 20 + assert reset_state["idle_pressure_score"] == 0 + + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(5, 16 << 20, 8 * gib) + jump_state = self._pool_state(dc) + assert jump_state["retries"] == 5 + assert jump_state["budget"] == 4 << 20 + assert jump_state["idle_pressure_score"] == 5 + finally: + dc.cleanup() + + def test_budget_lowering_discards_checked_out_storage_on_return(self): + graph_id, first_ds_id, checked_out_ds_id = 9110, 9111, 9112 + dc = self._init_dc(pool_budget=20_000) + try: + first_shard = self._register_param(dc, graph_id, first_ds_id, [4097], register_graph=False) + checked_out_shard = self._register_param(dc, graph_id, checked_out_ds_id, [2049], register_graph=False) + dc.register_graph_z3(graph_id, [first_ds_id, checked_out_ds_id]) + + first_view, first_storage = self._gather_view_and_storage(first_shard, graph_id, first_ds_id) + self._release(first_view, graph_id, first_ds_id, 1) + checked_out_view, checked_out_storage = self._gather_view_and_storage(checked_out_shard, graph_id, + checked_out_ds_id) + assert checked_out_storage.data_ptr() == first_storage.data_ptr() + + dc.set_z3_gather_buffer_pool_budget_for_test(0) + self._release(checked_out_view, graph_id, checked_out_ds_id, 1) + assert checked_out_storage.nbytes() == 0 + finally: + dc.cleanup() + + def test_cleanup_clears_process_global_test_override(self): + first_graph_id, first_ds_id = 9120, 9121 + dc = self._init_dc() + first_shard = self._register_param(dc, first_graph_id, first_ds_id, [4097]) + first_view, first_storage = self._gather_view_and_storage(first_shard, first_graph_id, first_ds_id) + self._release(first_view, first_graph_id, first_ds_id, 1) + assert first_storage.nbytes() > 0 + dc.cleanup() + + second_graph_id, second_ds_id = 9122, 9123 + dc = self._init_dc(pool_budget=None) + try: + second_shard = self._register_param(dc, second_graph_id, second_ds_id, [4097]) + second_view, second_storage = self._gather_view_and_storage(second_shard, second_graph_id, second_ds_id) + self._release(second_view, second_graph_id, second_ds_id, 1) + assert second_storage.nbytes() == 0 + finally: + dc.cleanup() + + def test_cleanup_is_idempotent_and_releases_registry_state(self): + graph_id, ds_id = 9160, 9161 + dc = self._init_dc() + first_shard = self._register_param(dc, graph_id, ds_id, [4097]) + first_view, _ = self._gather_view_and_storage(first_shard, graph_id, ds_id) + self._release(first_view, graph_id, ds_id, 1) + + dc.cleanup() + dc.cleanup() + + dc = self._init_dc() + try: + second_shard = self._register_param(dc, graph_id, ds_id, [2049]) + second_view, _ = self._gather_view_and_storage(second_shard, graph_id, ds_id) + assert torch.allclose(second_view.sum(), self._expected_view_sum([2049])) + self._release(second_view, graph_id, ds_id, 1) + finally: + dc.cleanup() + + def test_profile_invalidation_discards_checked_out_pool_storage_immediately(self): + graph_id = 9080 + first_ds_id, invalidated_ds_id, next_ds_id, reused_ds_id = 9081, 9082, 9083, 9084 + dc = self._init_dc() + try: + first_shard = self._register_param(dc, graph_id, first_ds_id, [4097], register_graph=False) + invalidated_shard = self._register_param(dc, graph_id, invalidated_ds_id, [2049], register_graph=False) + next_shard = self._register_param(dc, graph_id, next_ds_id, [1025], register_graph=False) + reused_shard = self._register_param(dc, graph_id, reused_ds_id, [513], register_graph=False) + dc.register_graph_z3(graph_id, [first_ds_id, invalidated_ds_id, next_ds_id, reused_ds_id]) + + first_view, first_storage = self._gather_view_and_storage(first_shard, graph_id, first_ds_id) + pool_ptr = first_storage.data_ptr() + self._release(first_view, graph_id, first_ds_id, 1) + + invalidated_view, invalidated_storage = self._gather_view_and_storage(invalidated_shard, graph_id, + invalidated_ds_id) + assert invalidated_storage.data_ptr() == pool_ptr + dc.invalidate_gathered_param(invalidated_ds_id) + + # The invalidated view deliberately remains alive here. A new gather must + # not acquire the storage that profiling removed from pool ownership. + next_view, next_storage = self._gather_view_and_storage(next_shard, graph_id, next_ds_id) + assert next_storage.data_ptr() != pool_ptr + + self._release(next_view, graph_id, next_ds_id, 1) + dc.clear_all_gathered_params() + del invalidated_view + get_accelerator().synchronize() + + reused_view, reused_storage = self._gather_view_and_storage(reused_shard, graph_id, reused_ds_id) + assert reused_storage.data_ptr() == next_storage.data_ptr() + self._release(reused_view, graph_id, reused_ds_id, 1) + finally: + dc.cleanup() + def test_persistent_param_storage_unchanged_across_release(self): graph_id, ds_id = 9030, 9031 dc = self._init_dc() @@ -103,14 +1058,16 @@ def test_persistent_param_storage_unchanged_across_release(self): dc.cleanup() def test_consumer_stream_can_finish_before_storage_reuse(self): - graph_id, ds_id = 9040, 9041 + graph_id, ds_id, next_ds_id = 9040, 9041, 9042 if not hasattr(torch.cuda, "_sleep"): #ignore-cuda pytest.skip("CUDA sleep helper is unavailable") dc = self._init_dc() try: - shard = self._register_param(dc, graph_id, ds_id, [4097]) + shard = self._register_param(dc, graph_id, ds_id, [4097], register_graph=False) + next_shard = self._register_param(dc, graph_id, next_ds_id, [2049], register_graph=False) + dc.register_graph_z3(graph_id, [ds_id, next_ds_id]) view, storage = self._gather_view_and_storage(shard, graph_id, ds_id) - padded_bytes = storage.nbytes() + before_ptr = storage.data_ptr() result = torch.empty((), device=self._device(), dtype=view.dtype) consumer_stream = get_accelerator().Stream() with get_accelerator().stream(consumer_stream): @@ -118,13 +1075,46 @@ def test_consumer_stream_can_finish_before_storage_reuse(self): result.copy_(view.sum()) self._release(view, graph_id, ds_id, 1, synchronize=False) - scratch = torch.empty((padded_bytes // view.element_size()) + 1024, - device=self._device(), - dtype=view.dtype) - scratch.fill_(17) + next_view, next_storage = self._gather_view_and_storage(next_shard, graph_id, next_ds_id) get_accelerator().synchronize() assert torch.allclose(result, self._expected_view_sum([4097])) - assert storage.nbytes() == 0 - del scratch + assert next_storage.data_ptr() == before_ptr + assert torch.allclose(next_view.sum(), self._expected_view_sum([2049])) + assert storage.nbytes() > 0 + self._release(next_view, graph_id, next_ds_id, 1) + finally: + dc.cleanup() + + def test_pressure_recovery_waits_for_recorded_consumer_stream_before_reallocation(self): + graph_id, ds_id = 9179, 9180 + if not hasattr(torch.cuda, "_sleep"): #ignore-cuda + pytest.skip("CUDA sleep helper is unavailable") + dc = self._init_dc(pool_budget=None) + try: + shard = self._register_param(dc, graph_id, ds_id, [1_048_577]) + gib = 1 << 30 + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(0, gib, 8 * gib) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(1, gib, 8 * gib) + + first_view, _ = self._gather_view_and_storage(shard, graph_id, ds_id) + self._release(first_view, graph_id, ds_id, 1) + held_view, _ = self._gather_view_and_storage(shard, graph_id, ds_id) + dc.update_z3_gather_buffer_pool_allocator_pressure_for_test(5, 8 << 20, 8 * gib) + + result = torch.empty((), device=self._device(), dtype=held_view.dtype) + consumer_stream = get_accelerator().Stream() + with get_accelerator().stream(consumer_stream): + torch.cuda._sleep(int(1e8)) #ignore-cuda + result.copy_(held_view.sum()) + self._release(held_view, graph_id, ds_id, 1, synchronize=False) + + # The next acquire performs the one-shot allocator flush. The + # recorded consumer stream must finish before its retired block can + # participate in the new gather allocation. + next_view, _ = self._gather_view_and_storage(shard, graph_id, ds_id) + get_accelerator().synchronize() + assert torch.allclose(result, self._expected_view_sum([1_048_577])) + assert torch.allclose(next_view.sum(), self._expected_view_sum([1_048_577])) + self._release(next_view, graph_id, ds_id, 1) finally: dc.cleanup() diff --git a/tests/unit/compile/test_backend.py b/tests/unit/compile/test_backend.py new file mode 100644 index 000000000000..c848c992d5b9 --- /dev/null +++ b/tests/unit/compile/test_backend.py @@ -0,0 +1,359 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from collections import deque +from contextlib import nullcontext +from types import MethodType, SimpleNamespace + +import torch + +from deepspeed.compile.backend import _get_fw_real_inputs, set_example_values_to_symints +from deepspeed.compile import backend as backend_mod +from deepspeed.compile.inductor import patch_create_aot_dispatcher_function +from deepspeed.compile.graph_param import DSGraphParamManager +from deepspeed.compile.input_storage import InputStorage +from deepspeed.compile.patch_compiled_func import (clear_backward_inputs, get_backward_inputs, patch_compiled_func, + unpatch_compiled_func) +from deepspeed.compile.profilers import ProfilingResult +from deepspeed.compile.profilers import graph_profile as graph_profile_mod +from deepspeed.compile.profilers.graph_profile import _mark_profile_incomplete + +_DC_LIBRARIES = [] + + +def test_profiling_result_keeps_existing_positional_field_order(): + graph = torch.fx.Graph() + + result = ProfilingResult(graph) + + assert result.fwd_graph is graph + assert result.process_group is None + + +def _define_dc_ops(): + try: + torch.ops.dc.allgather_param.default + torch.ops.dc.wait_allgather.default + torch.ops.dc.release_param.default + torch.ops.dc.reduce_grad.default + return + except AttributeError: + pass + + lib = torch.library.Library("dc", "FRAGMENT") + for schema in ( + "allgather_param(Tensor a, int graph_id, int id, ScalarType? dtype = None) -> Tensor", + "wait_allgather(Tensor(a) a, int graph_id, int id) -> Tensor(a)", + "release_param(Tensor(a) a, int graph_id, int id, int n_users) -> Tensor(a)", + "reduce_grad(Tensor a, int graph_id, int id) -> Tensor", + "free_tensors(Tensor[] tensors) -> ()", + "end_backward(Tensor[] tensors, int graph_id, bool release_reduce_buckets = True) -> ()", + ): + try: + lib.define(schema) + except RuntimeError as exc: + if "already been registered" not in str(exc): + raise + _DC_LIBRARIES.append(lib) + + +def test_forward_real_inputs_are_graph_local(): + local_inputs = (torch.nn.Parameter(torch.ones(2, dtype=torch.float32)), ) + storage = InputStorage() + storage.put((torch.ones(1, dtype=torch.float32), )) + + selected = _get_fw_real_inputs(deque([local_inputs]), storage, graph_id=7) + + assert selected is local_inputs + + +def test_forward_real_inputs_fall_back_to_storage_when_local_queue_is_empty(): + storage = InputStorage() + storage.put((torch.ones(3, dtype=torch.float32), )) + + selected = _get_fw_real_inputs(deque(), storage, graph_id=7) + + assert len(selected) == 1 + assert selected[0].shape == torch.Size([3]) + assert selected[0].dtype is torch.float32 + + +def test_stored_zero_parameter_recovers_original_instance_bound_protocol(): + real_param = torch.nn.Parameter(torch.ones(3), requires_grad=False) + real_param.ds_id = 123 + real_param.ds_shape = torch.Size([3]) + real_param.ds_persist = False + real_param.all_gather = MethodType(lambda self, param_list: None, real_param) + real_param.partition = MethodType(lambda self, param_list, has_been_updated: None, real_param) + storage = InputStorage() + storage.put((real_param, )) + + stored_inputs = storage.get() + materialized = set_example_values_to_symints(stored_inputs, [(0, 123, torch.Size([3]))], + real_zero_params={123: real_param}) + + assert stored_inputs[0] is not real_param + assert not hasattr(stored_inputs[0], "ds_id") + assert materialized[0] is real_param + assert materialized[0].all_gather.__self__ is real_param + assert materialized[0].partition.__self__ is real_param + + +def test_symint_materialization_preserves_frozen_zero_parameter_for_profiling_consumers(monkeypatch): + from torch._subclasses.fake_tensor import FakeTensorMode + + _define_dc_ops() + calls = [] + real_param = torch.nn.Parameter(torch.empty((2, 3), dtype=torch.bfloat16), requires_grad=False) + real_param.ds_id = 123 + real_param.ds_shape = torch.Size([2, 3]) + real_param.ds_persist = False + + def all_gather(self, param_list): + calls.append(("all_gather", param_list)) + + def partition(self, param_list, has_been_updated): + calls.append(("partition", param_list, has_been_updated)) + + real_param.all_gather = MethodType(all_gather, real_param) + real_param.partition = MethodType(partition, real_param) + + with FakeTensorMode() as fake_mode: + fake_param = fake_mode.from_tensor(real_param) + fake_param.ds_id = 123 + fake_param.ds_shape = torch.Size([2, 3]) + fake_param.ds_persist = False + + materialized = set_example_values_to_symints((fake_param, ), [(0, 123, torch.Size([2, 3]))], + real_zero_params={123: real_param}) + assert materialized[0] is real_param + graph = torch.fx.Graph() + param_node = graph.placeholder("frozen_zero_param") + neg_node = graph.call_function(torch.neg, (param_node, )) + graph.output((neg_node, )) + manager = DSGraphParamManager(graph, materialized, [(0, 123, torch.Size([2, 3]))]) + managed_param = manager.params[param_node.name].param + persistent_ds_ids = { + manager.ds_ids[name] + for name, graph_param in manager.params.items() if graph_param.param.ds_persist + } + + assert isinstance(managed_param, torch.nn.Parameter) + assert managed_param.shape == torch.Size([2, 3]) + assert managed_param.dtype is torch.bfloat16 + assert not managed_param.requires_grad + assert managed_param.ds_id == 123 + assert managed_param.ds_shape == torch.Size([2, 3]) + assert managed_param.ds_persist is False + assert persistent_ds_ids == set() + + class FakeEvent: + + def record(self): + pass + + def elapsed_time(self, other): + return 0.0 + + class FakeAccelerator: + + def current_device(self): + return "cpu" + + def random(self): + return SimpleNamespace(fork_rng=lambda devices: nullcontext()) + + def Event(self, enable_timing): + return FakeEvent() + + def reset_peak_memory_stats(self): + pass + + def memory_allocated(self): + return 100 + + def max_memory_allocated(self): + return 100 + + def synchronize(self): + pass + + class FakeDeepCompileHandle: + + def enable_profiling(self, enabled): + pass + + def clear_all_gathered_params(self): + pass + + monkeypatch.setattr(graph_profile_mod, "get_accelerator", lambda: FakeAccelerator()) + monkeypatch.setattr(graph_profile_mod, "get_deepcompile_handle", lambda: FakeDeepCompileHandle()) + monkeypatch.setattr(graph_profile_mod, "_get_mem_usage_out_of_torch", lambda: 0) + monkeypatch.setattr(graph_profile_mod, "is_comm_op", lambda node: False) + monkeypatch.setattr(graph_profile_mod, "is_release_node", lambda node: False) + monkeypatch.setattr(graph_profile_mod.dist, "is_initialized", lambda: False) + monkeypatch.setattr(graph_profile_mod.dist, "get_rank", lambda: 0) + + gm = torch.fx.GraphModule(torch.nn.Module(), graph) + profiler = graph_profile_mod.ProfilingInterpreter(gm, iteration=1, warmup=0) + profiler.run(*materialized) + + assert not graph_profile_mod.is_profile_incomplete(graph) + assert [call[0] for call in calls] == ["all_gather", "partition"] + + +def test_launch_compile_passes_clears_owned_compiled_backward_state(monkeypatch): + + class DummyDeepCompileHandle: + + def reset(self): + pass + + clear_backward_inputs() + backend_mod.frames_needing_bwd.clear() + unpatch_compiled_func() + original_autograd_function = torch.autograd.Function + owner = object() + owned_frames = {(owner, 17)} + backend_mod.frames_needing_bwd.update(owned_frames) + patch_compiled_func() + get_backward_inputs().append((torch.ones(1), )) + monkeypatch.setattr(backend_mod, "log_rank0", lambda *args, **kwargs: None) + monkeypatch.setattr(backend_mod, "get_deepcompile_handle", lambda: DummyDeepCompileHandle()) + + backend_mod.init_schedule([(0, [])]) + try: + backend_mod.launch_compile_passes(0, owned_frames=owned_frames) + + assert owned_frames == set() + assert backend_mod.frames_needing_bwd == set() + assert get_backward_inputs() == [] + assert torch.autograd.Function is original_autograd_function + finally: + backend_mod.frames_needing_bwd.clear() + unpatch_compiled_func() + + +def test_unpatch_compiled_func_clears_backward_inputs(): + clear_backward_inputs() + patch_compiled_func() + try: + get_backward_inputs().append((torch.ones(1), )) + unpatch_compiled_func() + assert get_backward_inputs() == [] + finally: + unpatch_compiled_func() + + +def _patch_aot_constructor(): + return patch_create_aot_dispatcher_function(graph_id=7, + z3_partition=False, + make_fw_graph=lambda gm, sample_inputs: gm.graph, + make_bw_graph=lambda gm, sample_inputs: gm.graph, + real_inputs=(torch.ones(1), ), + param_indices=[], + param_manager={}, + frame_id=0, + frames_partitioned=set()) + + +def test_inductor_aot_constructor_patch_is_restorable(): + from torch._dynamo.backends.common import AotAutograd + + original_init = AotAutograd.__init__ + restore = patch_create_aot_dispatcher_function(graph_id=7, + z3_partition=False, + make_fw_graph=lambda gm, sample_inputs: gm.graph, + make_bw_graph=lambda gm, sample_inputs: gm.graph, + real_inputs=(torch.ones(1), ), + param_indices=[], + param_manager={}, + frame_id=0, + frames_partitioned=set()) + try: + assert AotAutograd.__init__ is not original_init + finally: + restore() + + assert AotAutograd.__init__ is original_init + assert not hasattr(AotAutograd, "__original_init") + + +def test_older_aot_restore_does_not_clobber_newer_patch(): + from torch._dynamo.backends.common import AotAutograd + + original_init = AotAutograd.__init__ + restore_first = _patch_aot_constructor() + restore_second = _patch_aot_constructor() + newer_init = AotAutograd.__init__ + try: + restore_first() + assert AotAutograd.__init__ is newer_init + assert hasattr(AotAutograd, "__original_init") + finally: + restore_second() + + assert AotAutograd.__init__ is original_init + assert not hasattr(AotAutograd, "__original_init") + + +def test_run_opt_passes_skips_memory_profile_for_incomplete_graph(monkeypatch): + gm = torch.fx.symbolic_trace(lambda x: x + 1) + profiling_results = {7: ProfilingResult()} + + class UnexpectedMemoryProfiler: + + def __init__(self, *args, **kwargs): + raise AssertionError("memory profiling should be skipped for incomplete operator profiles") + + def incomplete_profile_pass(gm, *args, **kwargs): + _mark_profile_incomplete(gm.graph) + return gm + + monkeypatch.setattr(backend_mod, "MemoryProfilingInterpreter", UnexpectedMemoryProfiler) + monkeypatch.setattr(backend_mod, "log_rank0", lambda *args, **kwargs: None) + + backend_mod.run_opt_passes(opt_passes=[incomplete_profile_pass], + gm=gm, + graph_id=7, + graph_order=[], + profiling_results=profiling_results, + create_inputs_fn=lambda: (torch.ones(1), ), + mem_budget=0.0, + param_manager={}, + bwd=False) + + assert profiling_results[7].fwd_mem == [] + assert profiling_results[7].fwd_mem_complete is False + + +def test_run_opt_passes_skips_memory_profile_when_another_rank_is_incomplete(monkeypatch): + gm = torch.fx.symbolic_trace(lambda x: x + 1) + profiling_results = {7: ProfilingResult()} + + class UnexpectedMemoryProfiler: + + def __init__(self, *args, **kwargs): + raise AssertionError("all ranks must skip profiling when any rank has an incomplete operator profile") + + def complete_profile_pass(gm, *args, **kwargs): + return gm + + monkeypatch.setattr(backend_mod, "MemoryProfilingInterpreter", UnexpectedMemoryProfiler) + monkeypatch.setattr(backend_mod, "_sync_memory_profile_complete", lambda complete, process_group=None: False) + monkeypatch.setattr(backend_mod, "log_rank0", lambda *args, **kwargs: None) + + backend_mod.run_opt_passes(opt_passes=[complete_profile_pass], + gm=gm, + graph_id=7, + graph_order=[], + profiling_results=profiling_results, + create_inputs_fn=lambda: (torch.ones(1), ), + mem_budget=0.0, + param_manager={}, + bwd=False) + + assert profiling_results[7].fwd_mem == [] + assert profiling_results[7].fwd_mem_complete is False diff --git a/tests/unit/compile/test_inductor_aot_kwargs.py b/tests/unit/compile/test_inductor_aot_kwargs.py index d2784fc09927..f8b9437a3cbf 100644 --- a/tests/unit/compile/test_inductor_aot_kwargs.py +++ b/tests/unit/compile/test_inductor_aot_kwargs.py @@ -3,6 +3,8 @@ # DeepSpeed Team +import torch + import deepspeed.compile.inductor as inductor @@ -139,3 +141,97 @@ def test_torchxla_openxla_shape_passes_through_unchanged(monkeypatch): assert result["compiler_calls"] == [] assert result["partition_calls"] == [] assert kwargs == original_kwargs + + +def test_deepcompile_z3_inductor_config_patch_disables_available_reduction_heuristics(): + config = torch._inductor.config + triton_config = config.triton + original_values = { + config_name: getattr(triton_config, + config_name.split(".", 1)[1]) + for config_name in inductor._DEEP_COMPILE_Z3_INDUCTOR_REDUCTION_CONFIG + if hasattr(triton_config, + config_name.split(".", 1)[1]) + } + assert original_values + + with inductor.deepcompile_z3_inductor_config_patch(enabled=True): + for config_name in original_values: + assert getattr(triton_config, config_name.split(".", 1)[1]) is False + + for config_name, original_value in original_values.items(): + assert getattr(triton_config, config_name.split(".", 1)[1]) is original_value + + with inductor.deepcompile_z3_inductor_config_patch(enabled=False): + for config_name, original_value in original_values.items(): + assert getattr(triton_config, config_name.split(".", 1)[1]) is original_value + + +def test_deepcompile_z3_inductor_config_patch_skips_unavailable_reduction_heuristics(monkeypatch): + + class FakeTritonConfig: + persistent_reductions = True + + class FakePatch: + + def __init__(self, overrides): + self.overrides = overrides + + def __enter__(self): + assert self.overrides == {"triton.persistent_reductions": False} + FakeTritonConfig.persistent_reductions = False + + def __exit__(self, exc_type, exc, tb): + FakeTritonConfig.persistent_reductions = True + + class FakeConfig: + triton = FakeTritonConfig() + + @staticmethod + def patch(overrides): + return FakePatch(overrides) + + class FakeInductor: + config = FakeConfig() + + monkeypatch.setattr(torch, "_inductor", FakeInductor()) + + with inductor.deepcompile_z3_inductor_config_patch(enabled=True): + assert FakeTritonConfig.persistent_reductions is False + + assert FakeTritonConfig.persistent_reductions is True + + +def test_patch_compiler_applies_z3_inductor_config_during_original_compile(monkeypatch): + events = [] + + class ConfigContext: + + def __init__(self, enabled): + self.enabled = enabled + + def __enter__(self): + events.append(("enter", self.enabled)) + + def __exit__(self, exc_type, exc, tb): + events.append(("exit", self.enabled)) + + monkeypatch.setattr(inductor, "deepcompile_z3_inductor_config_patch", lambda enabled: ConfigContext(enabled)) + + class ParamManager: + param_names = [] + + def original_compiler(gm, inputs): + events.append(("compile", tuple(inputs))) + return "compiled" + + gm = torch.fx.symbolic_trace(lambda: torch.ones(())) + wrapped = inductor.patch_compiler(original_compiler, + dc_compiler=lambda gm, inputs: gm.graph, + z3_partition=True, + graph_id=7, + graph_param_manager={7: ParamManager()}, + bwd=False) + + assert wrapped(gm, ()) == "compiled" + assert events == [("enter", True), ("compile", ()), ("exit", True)] diff --git a/tests/unit/compile/test_list_schedule.py b/tests/unit/compile/test_list_schedule.py index 873ea989761f..22a08bc7592c 100644 --- a/tests/unit/compile/test_list_schedule.py +++ b/tests/unit/compile/test_list_schedule.py @@ -4,6 +4,7 @@ # DeepSpeed Team import operator +import sys from types import SimpleNamespace import pytest @@ -16,7 +17,9 @@ from deepspeed.compile import list_schedule as schedule_mod from deepspeed.compile.passes import prefetch as prefetch_mod from deepspeed.compile.passes import selective_gather as selective_gather_mod +from deepspeed.compile.passes import zero3_compile as zero3_compile_mod from deepspeed.compile.profilers import ProfilingResult +from deepspeed.compile.profilers import graph_profile as graph_profile_mod from deepspeed.compile.profilers.graph_profile import _backfill_missing_profile_metadata, is_profile_incomplete _DC_LIBRARIES = [] @@ -58,6 +61,9 @@ def stub_deepcompile_ops(monkeypatch): def _with_meta(node, tensor_size=0, device_time=0): node.meta["tensor_size"] = tensor_size + node.meta["alloc_mem"] = 0 + node.meta["profile_mem_start"] = 0 + node.meta["profile_mem_peak"] = 0 if device_time is not None: node.meta["device_time"] = device_time return node @@ -82,23 +88,870 @@ def fail_all_reduce(*args, **kwargs): def test_sync_memory_profile_complete_reduces_asymmetric_failure(monkeypatch): monkeypatch.setattr(backend_mod.dist, "is_initialized", lambda: True) monkeypatch.setattr(backend_mod, "get_accelerator", lambda: SimpleNamespace(current_device=lambda: "cpu")) + process_group = object() - def mark_any_rank_failed(tensor, op): + def mark_any_rank_failed(tensor, op, group=None): assert op == backend_mod.dist.ReduceOp.MIN + assert group is process_group tensor[0] = 0 monkeypatch.setattr(backend_mod.dist, "all_reduce", mark_any_rank_failed) - assert not backend_mod._sync_memory_profile_complete(True) + assert not backend_mod._sync_memory_profile_complete(True, process_group) -def _allgather(graph, arg, ds_id, name, tensor_size=1, device_time=1): - return _with_meta( +def test_get_last_uses_handles_dead_no_copy_node(): + graph = Graph() + param = _placeholder(graph, "dead_wait_param") + wait = _wait(graph, param, 1, "dead_wait") + graph.output(()) + graph.lint() + + node_to_last_use, user_to_last_uses = compile_util.get_last_uses(graph) + node_to_uses = compile_util.get_real_uses(graph) + + assert node_to_last_use[param] is wait + assert user_to_last_uses[wait] == [param] + assert node_to_uses[param] == [] + + +def test_zero3_scheduler_budget_uses_rank_reduced_non_gathered_peak(monkeypatch): + monkeypatch.setattr(zero3_compile_mod.dist, "is_initialized", lambda: True) + + class FakeAccelerator: + + def current_device(self): + return "cpu" + + def total_memory(self): + return 2000 + + def memory_allocated(self): + return 50 + + max_reductions = iter((850, 200)) + + def reduce_budget_inputs(tensor, op): + if op == zero3_compile_mod.dist.ReduceOp.MIN: + tensor[0] = 1000 + elif op == zero3_compile_mod.dist.ReduceOp.MAX: + tensor[0] = next(max_reductions) + else: + raise AssertionError(f"unexpected reduce op {op}") + + monkeypatch.setattr(zero3_compile_mod, "get_accelerator", lambda: FakeAccelerator()) + monkeypatch.setattr(zero3_compile_mod.dist, "all_reduce", reduce_budget_inputs) + + graph = Graph() + param = _placeholder(graph, "budget_builder_param") + ag = _allgather(graph, param, 1, "budget_builder", tensor_size=200) + wait = _wait(graph, ag, 1, "budget_builder") + op = _neg(graph, wait, "budget_builder_op") + op.meta.update(max_mem=800, profile_mem_start=250, profile_mem_peak=1050) + release = _release(graph, op, 1, "budget_builder") + graph.output((release, )) + graph.lint() + for node in graph.nodes: + node.meta.setdefault("alloc_mem", 0) + node.meta.setdefault("max_mem", 0) + node.meta.setdefault("profile_mem_start", 0) + node.meta.setdefault("profile_mem_peak", 0) + + budget = zero3_compile_mod._build_scheduler_budget_from_operator_profile(graph) + + assert budget.source == "profiled_non_gathered_peak_memory_clamped_to_minimum_gather_residency" + assert budget.total_mem == 1000 + assert budget.profiled_non_gathered_peak_mem == 850 + assert budget.safety_margin == 100 + assert budget.available_mem == 50 + assert budget.max_gathered_bytes == 200 + + +def test_profiled_scheduler_budget_clamps_impossible_step7_shape_to_largest_single_gather(): + budget = schedule_mod.SchedulerMemoryBudget.from_profiled_non_gathered_peak( + total_mem=85017493504, + profiled_non_gathered_peak_mem=75524405760, + output_size=0, + ) + + clamped = budget.clamped_to_minimum_gather_residency(1555824640) + + assert budget.max_gathered_bytes == 991338394 + assert clamped.max_gathered_bytes == 1555824640 + assert clamped.available_mem == 991338394 + assert clamped.profiled_non_gathered_peak_mem == 75524405760 + assert clamped.source == "profiled_non_gathered_peak_memory_clamped_to_minimum_gather_residency" + + +def test_zero3_scheduler_complete_profile_clamps_to_largest_single_gather(monkeypatch): + monkeypatch.setattr(zero3_compile_mod.dist, "is_initialized", lambda: True) + monkeypatch.setattr(zero3_compile_mod.dist, "get_world_size", lambda group=None: 1) + monkeypatch.setattr(zero3_compile_mod.dist, "all_reduce", lambda *args, **kwargs: None) + + class FakeAccelerator: + + def current_device(self): + return "cpu" + + def total_memory(self): + return 1000 + + monkeypatch.setattr(zero3_compile_mod, "get_accelerator", lambda: FakeAccelerator()) + + graph = Graph() + first_param = _placeholder(graph, "profile_floor_first_param") + second_param = _placeholder(graph, "profile_floor_second_param") + first_ag = _allgather(graph, first_param, 1, "profile_floor_first", tensor_size=800, allocation_size=800) + first_wait = _wait(graph, first_ag, 1, "profile_floor_first") + second_ag = _allgather(graph, second_param, 2, "profile_floor_second", tensor_size=800, allocation_size=800) + second_wait = _wait(graph, second_ag, 2, "profile_floor_second") + op = _add(graph, first_wait, second_wait, "profile_floor_op") + op.meta.update(max_mem=800, profile_mem_start=850, profile_mem_peak=850) + first_release = _release(graph, op, 1, "profile_floor_first") + second_release = _release(graph, first_release, 2, "profile_floor_second") + graph.output((second_release, )) + graph.lint() + for node in graph.nodes: + node.meta.setdefault("alloc_mem", 0) + node.meta.setdefault("max_mem", 0) + node.meta.setdefault("profile_mem_start", 0) + node.meta.setdefault("profile_mem_peak", 0) + + budget, disabled_reason = zero3_compile_mod._scheduler_budget_from_operator_profile( + GraphModule(torch.nn.Module(), graph)) + + assert budget.max_gathered_bytes == 800 + assert budget.available_mem == 50 + assert budget.source == "profiled_non_gathered_peak_memory_clamped_to_minimum_gather_residency" + assert disabled_reason is None + + +def test_zero3_scheduler_budget_reconstructs_live_activations_before_transient_peak(monkeypatch): + monkeypatch.setattr(zero3_compile_mod.dist, "is_initialized", lambda: False) + monkeypatch.setattr(zero3_compile_mod, "get_accelerator", lambda: SimpleNamespace(memory_allocated=lambda: 100)) + + graph = Graph() + value = _placeholder(graph, "absolute_peak_input") + activation = _neg(graph, value, "absolute_peak_activation") + activation.meta.update(alloc_mem=300, max_mem=300, profile_mem_start=100, profile_mem_peak=400) + transient = _neg(graph, activation, "absolute_peak_transient") + transient.meta.update(alloc_mem=50, max_mem=200, profile_mem_start=400, profile_mem_peak=600) + graph.output((transient, )) + graph.lint() + + peak = zero3_compile_mod._rank_max_operator_profiled_non_gathered_peak(graph) + + # The absolute record retains the first node's live activation while the + # second node reaches a further transient peak. + assert peak == 600 + + +def test_zero3_scheduler_budget_uses_absolute_peaks_after_inter_node_reclamation(monkeypatch): + monkeypatch.setattr(zero3_compile_mod.dist, "is_initialized", lambda: False) + + graph = Graph() + value = _placeholder(graph, "reclaimed_peak_input") + first = _neg(graph, value, "reclaimed_first") + first.meta.update(alloc_mem=400, max_mem=400, profile_mem_start=100, profile_mem_peak=500) + second = _neg(graph, first, "reclaimed_second") + second.meta.update(alloc_mem=-300, max_mem=250, profile_mem_start=200, profile_mem_peak=450) + graph.output((second, )) + graph.lint() + + peak = zero3_compile_mod._rank_max_operator_profiled_non_gathered_peak(graph) + + assert peak == 500 + + +def test_absolute_profile_memory_adds_external_once_and_max_reduces_asymmetric_ranks(monkeypatch): + + class FakeAccelerator: + + def memory_allocated(self): + return 100 + + def max_memory_allocated(self): + return 140 + + monkeypatch.setattr(graph_profile_mod, "get_accelerator", lambda: FakeAccelerator()) + + assert graph_profile_mod._absolute_profile_memory(50) == (150, 190) + + def reduce_to_worst_rank(values, op): + assert op == graph_profile_mod.dist.ReduceOp.MAX + values[0] = 175 + values[1] = 410 + + monkeypatch.setattr(graph_profile_mod.dist, "all_reduce", reduce_to_worst_rank) + assert graph_profile_mod._rank_max_profile_memory(150, 190, torch.device("cpu"), distributed=True) == (175, 410) + + +def test_external_memory_profile_maps_nvml_device_and_excludes_reserved_allocator_bytes(monkeypatch): + nvml_device_ids = [] + fake_nvml = SimpleNamespace( + nvmlInit=lambda: None, + nvmlDeviceGetHandleByIndex=lambda device_id: nvml_device_ids.append(device_id) or "handle", + nvmlDeviceGetMemoryInfo=lambda handle: SimpleNamespace(used=1000)) + monkeypatch.setitem(sys.modules, "pynvml", fake_nvml) + + class FakeAccelerator: + + def current_device(self): + return 0 + + def _get_nvml_gpu_id(self, device_id): + assert device_id == 0 + return 5 + + def memory_allocated(self): + return 400 + + def memory_reserved(self): + return 600 + + monkeypatch.setattr(graph_profile_mod, "get_accelerator", lambda: FakeAccelerator()) + + assert graph_profile_mod._get_mem_usage_out_of_torch() == 400 + assert nvml_device_ids == [5] + + +def test_external_memory_profile_does_not_count_reserved_to_allocated_reuse_twice(monkeypatch): + fake_nvml = SimpleNamespace(nvmlInit=lambda: None, + nvmlDeviceGetHandleByIndex=lambda device_id: "handle", + nvmlDeviceGetMemoryInfo=lambda handle: SimpleNamespace(used=1000)) + monkeypatch.setitem(sys.modules, "pynvml", fake_nvml) + + class FakeAccelerator: + allocated = 400 + + def current_device(self): + return 0 + + def memory_allocated(self): + return self.allocated + + def max_memory_allocated(self): + return self.allocated + + def memory_reserved(self): + return 800 + + accelerator = FakeAccelerator() + monkeypatch.setattr(graph_profile_mod, "get_accelerator", lambda: accelerator) + + external_mem = graph_profile_mod._get_mem_usage_out_of_torch() + first_absolute, _ = graph_profile_mod._absolute_profile_memory(external_mem) + accelerator.allocated = 700 + second_absolute, _ = graph_profile_mod._absolute_profile_memory(external_mem) + + assert external_mem == 200 + assert (first_absolute, second_absolute) == (600, 900) + + +def test_zero3_scheduler_incomplete_profile_uses_identical_collective_on_asymmetric_ranks(monkeypatch): + monkeypatch.setattr(zero3_compile_mod.dist, "is_initialized", lambda: True) + monkeypatch.setattr(zero3_compile_mod.dist, "get_world_size", lambda: 2) + + class FakeAccelerator: + + def current_device(self): + return "cpu" + + def total_memory(self): + return 2000 + + def memory_allocated(self): + return 50 + + def available_memory(self): + return 500 + + monkeypatch.setattr(zero3_compile_mod, "get_accelerator", lambda: FakeAccelerator()) + + def make_graph(observed): + graph = Graph() + param = _placeholder(graph, f"asymmetric_param_{observed}") + ag = _allgather(graph, param, 1, f"asymmetric_{observed}", tensor_size=800) + wait = _wait(graph, ag, 1, f"asymmetric_{observed}") + op = _neg(graph, wait, f"asymmetric_op_{observed}") + if observed: + op.meta.update(alloc_mem=300, max_mem=500, profile_mem_start=850, profile_mem_peak=1300) + release = _release(graph, op, 1, f"asymmetric_{observed}") + graph.output((release, )) + graph.lint() + _backfill_missing_profile_metadata(graph, profile_complete=False) + return GraphModule(torch.nn.Module(), graph) + + sequences = [] + budgets = [] + for observed in (True, False): + calls = [] + + def reduce_asymmetric_rank(tensor, op): + calls.append((op, tensor.dtype)) + if op == zero3_compile_mod.dist.ReduceOp.MIN: + tensor[0] = 0 + elif op != zero3_compile_mod.dist.ReduceOp.MAX: + raise AssertionError(f"unexpected reduce op {op}") + + monkeypatch.setattr(zero3_compile_mod.dist, "all_reduce", reduce_asymmetric_rank) + budget, disabled_reason = zero3_compile_mod._scheduler_budget_from_operator_profile(make_graph(observed)) + assert budget.source == "incomplete_profile_minimum_gather_residency" + assert budget.max_gathered_bytes == 800 + assert disabled_reason is None + sequences.append(calls) + budgets.append(budget) + + assert sequences[0] == sequences[1] + assert sequences[0] == [(zero3_compile_mod.dist.ReduceOp.MIN, torch.int32), + (zero3_compile_mod.dist.ReduceOp.MAX, torch.int64)] + assert [budget.max_gathered_bytes for budget in budgets] == [800, 800] + + +def test_zero3_scheduler_budget_skips_incomplete_operator_profile_metadata(): + graph = Graph() + graph.output(()) + graph.lint() + _backfill_missing_profile_metadata(graph, profile_complete=False) + + budget = zero3_compile_mod._build_scheduler_budget_from_operator_profile(graph) + + assert budget is None + + +def test_zero3_scheduler_budget_skips_incomplete_operator_profile(monkeypatch): + monkeypatch.setattr(zero3_compile_mod.dist, "is_initialized", lambda: True) + monkeypatch.setattr(zero3_compile_mod, "get_accelerator", + lambda: SimpleNamespace(current_device=lambda: "cpu", available_memory=lambda: 0)) + monkeypatch.setattr(zero3_compile_mod.dist, "all_reduce", lambda *args, **kwargs: None) + graph = Graph() + graph.output(()) + graph.lint() + _backfill_missing_profile_metadata(graph, profile_complete=False) + gm = GraphModule(torch.nn.Module(), graph) + + budget, disabled_reason = zero3_compile_mod._scheduler_budget_from_operator_profile(gm) + + assert budget is None + assert disabled_reason == "incomplete_operator_profile_without_gather" + + +def test_zero3_scheduler_budget_caps_gathers_without_estimating_unprofiled_high_memory_suffix(monkeypatch): + monkeypatch.setattr(zero3_compile_mod.dist, "is_initialized", lambda: True) + monkeypatch.setattr(zero3_compile_mod.dist, "get_world_size", lambda: 1) + + class FakeAccelerator: + + def current_device(self): + return "cpu" + + def total_memory(self): + return 1000 + + def memory_allocated(self): + return 50 + + def available_memory(self): + return 500 + + def reduce_budget_inputs(tensor, op): + if tensor.dtype == torch.int32: + tensor[0] = 0 + + monkeypatch.setattr(zero3_compile_mod, "get_accelerator", lambda: FakeAccelerator()) + monkeypatch.setattr(zero3_compile_mod.dist, "all_reduce", reduce_budget_inputs) + + graph = Graph() + param = _placeholder(graph, "partial_budget_param") + ag = _allgather(graph, param, 1, "partial_budget", tensor_size=800) + wait = _wait(graph, ag, 1, "partial_budget") + op = _neg(graph, wait, "partial_budget_observed_op") + op.meta.update(max_mem=800, profile_mem_start=850, profile_mem_peak=1650) + unprofiled_suffix = _neg(graph, op, "partial_budget_unprofiled_high_memory_suffix") + release = _release(graph, op, 1, "partial_budget") + graph.output((release, unprofiled_suffix)) + graph.lint() + _backfill_missing_profile_metadata(graph, profile_complete=False) + gm = GraphModule(torch.nn.Module(), graph) + + budget, disabled_reason = zero3_compile_mod._scheduler_budget_from_operator_profile(gm) + + assert budget.source == "incomplete_profile_minimum_gather_residency" + assert budget.max_gathered_bytes == 800 + assert budget.available_mem == 0 + assert budget.profiled_non_gathered_peak_mem == 0 + assert disabled_reason is None + + +def test_minimum_gather_residency_budget_uses_largest_padded_gather_and_no_memory_estimate(monkeypatch): + monkeypatch.setattr(zero3_compile_mod.dist, "is_initialized", lambda: True) + monkeypatch.setattr(zero3_compile_mod.dist, "get_world_size", lambda group=None: 4) + monkeypatch.setattr(zero3_compile_mod, "get_accelerator", lambda: SimpleNamespace(current_device=lambda: "cpu")) + + def preserve_max(tensor, op, group=None): + assert op in (zero3_compile_mod.dist.ReduceOp.MIN, zero3_compile_mod.dist.ReduceOp.MAX) + + monkeypatch.setattr(zero3_compile_mod.dist, "all_reduce", preserve_max) + + graph = Graph() + small_param = _placeholder(graph, "minimum_residency_small_param") + large_param = _placeholder(graph, "minimum_residency_large_param") + small_ag = _allgather(graph, small_param, 1, "minimum_residency_small", tensor_size=10) + large_ag = _allgather(graph, large_param, 2, "minimum_residency_large", tensor_size=18) + graph.output((small_ag, large_ag)) + graph.lint() + _backfill_missing_profile_metadata(graph, profile_complete=False) + + budget, disabled_reason = zero3_compile_mod._scheduler_budget_from_operator_profile( + GraphModule(torch.nn.Module(), graph)) + + assert small_ag.meta["allgather_allocation_bytes"] == 16 + assert large_ag.meta["allgather_allocation_bytes"] == 24 + assert budget.max_gathered_bytes == 24 + assert budget.source == "incomplete_profile_minimum_gather_residency" + assert budget.available_mem == 0 + assert budget.total_mem == 0 + assert disabled_reason is None + + +def test_zero3_scheduler_budget_skips_non_distributed_memory_profile(monkeypatch): + monkeypatch.setattr(zero3_compile_mod.dist, "is_initialized", lambda: False) + graph = Graph() + graph.output(()) + graph.lint() + gm = GraphModule(torch.nn.Module(), graph) + + budget, disabled_reason = zero3_compile_mod._scheduler_budget_from_operator_profile(gm) + + assert budget is None + assert disabled_reason == "non_distributed" + + +def test_zero3_scheduler_budget_skips_when_budget_cannot_constrain(monkeypatch): + monkeypatch.setattr(zero3_compile_mod.dist, "is_initialized", lambda: True) + monkeypatch.setattr(zero3_compile_mod.dist, "get_world_size", lambda: 1) + monkeypatch.setattr(zero3_compile_mod.dist, "all_reduce", lambda *args, **kwargs: None) + + class FakeAccelerator: + + def current_device(self): + return "cpu" + + def total_memory(self): + return 2000 + + def memory_allocated(self): + return 50 + + monkeypatch.setattr(zero3_compile_mod, "get_accelerator", lambda: FakeAccelerator()) + + graph = Graph() + param = _placeholder(graph, "nonbinding_budget_param") + ag = _allgather(graph, param, 1, "nonbinding_budget", tensor_size=200) + wait = _wait(graph, ag, 1, "nonbinding_budget") + op = _neg(graph, wait, "nonbinding_budget_op") + op.meta.update(max_mem=800, profile_mem_start=250, profile_mem_peak=1050) + release = _release(graph, op, 1, "nonbinding_budget") + graph.output((release, )) + graph.lint() + for node in graph.nodes: + node.meta.setdefault("alloc_mem", 0) + node.meta.setdefault("max_mem", 0) + node.meta.setdefault("profile_mem_start", 0) + node.meta.setdefault("profile_mem_peak", 0) + gm = GraphModule(torch.nn.Module(), graph) + + budget, disabled_reason = zero3_compile_mod._scheduler_budget_from_operator_profile(gm) + + assert budget is None + assert disabled_reason == "budget_not_constraining" + + +def test_profiled_non_gathered_peak_conservatively_keeps_observed_peak(): + graph = Graph() + param = _placeholder(graph, "nongathered_peak_param") + ag = _allgather(graph, param, 2, "nongathered_peak", tensor_size=200) + wait = _wait(graph, ag, 2, "nongathered_peak") + release = _release(graph, wait, 2, "nongathered_peak") + graph.output((release, )) + graph.lint() + + assert schedule_mod.profiled_non_gathered_peak(graph, [(ag.name, 900, 0, 900), (wait.name, 950, 0, 950), + (release.name, 920, 0, 920)]) == 950 + + +def test_profiled_non_gathered_peak_does_not_subtract_nonresident_upfront_gathers(): + graph = Graph() + first_param = _placeholder(graph, "first_upfront_param") + second_param = _placeholder(graph, "second_upfront_param") + first_ag = _allgather(graph, first_param, 20, "first_upfront", tensor_size=200) + second_ag = _allgather(graph, second_param, 21, "second_upfront", tensor_size=300) + activation = _neg(graph, second_ag, "activation_heavy_node") + first_release = _release(graph, activation, 20, "first_upfront") + second_release = _release(graph, first_release, 21, "second_upfront") + graph.output((second_release, )) + graph.lint() + + # Operator profiling invalidates gathered buffers between nodes, so the + # activation peak is not guaranteed to include either upfront gather. The + # scheduler budget must not subtract hypothetical source-order residency. + mem_records = [(first_ag.name, 200, 0, 200), (second_ag.name, 300, 0, 300), (activation.name, 1000, 0, 1000)] + assert schedule_mod.profiled_non_gathered_peak(graph, mem_records) == 1000 + + +def test_zero3_stamps_padded_allgather_allocation_metadata(monkeypatch): + monkeypatch.setattr(zero3_compile_mod.dist, "is_initialized", lambda: True) + monkeypatch.setattr(zero3_compile_mod.dist, "get_world_size", lambda: 8) + graph = Graph() + + param = _placeholder(graph, "metadata_padded_param") + ag = _allgather(graph, param, 3, "metadata_padded", tensor_size=102) + wait = _wait(graph, ag, 3, "metadata_padded") + release = _release(graph, wait, 3, "metadata_padded") + graph.output((release, )) + graph.lint() + + zero3_compile_mod._set_allgather_allocation_metadata(graph) + + assert ag.meta["allgather_allocation_bytes"] == 112 + + +def test_zero3_stamps_replicated_param_allgather_allocation_metadata(): + graph = Graph() + + param = _placeholder(graph, "replicated_metadata_param") + param.meta["val"] = torch.empty((8, ), dtype=torch.float16) + use = _neg(graph, param, "replicated_metadata_use") + graph.output((use, )) + graph.lint() + + param_manager = SimpleNamespace(params={param.name: SimpleNamespace(dtype=torch.bfloat16, numel=777)}, + ds_ids={param.name: 3}) + + zero3_compile_mod.add_gather_and_release(0, graph, param_manager, [param]) + + ag_nodes = [node for node in graph.nodes if node.target == torch.ops.dc.allgather_param.default] + assert len(ag_nodes) == 1 + assert ag_nodes[0].meta["allgather_allocation_bytes"] == 1554 + + +def test_zero3_gathers_output_only_param_for_backward_passthrough(): + graph = Graph() + + param = _placeholder(graph, "output_only_param") + param.meta["val"] = torch.empty((8, ), dtype=torch.float16) + graph.output((param, )) + graph.lint() + + param_manager = SimpleNamespace(params={param.name: SimpleNamespace(dtype=torch.bfloat16, numel=777)}, + ds_ids={param.name: 3}) + + new_graph = zero3_compile_mod.add_gather_and_release(0, graph, param_manager, [param]) + new_graph.lint() + + ag_nodes = [node for node in new_graph.nodes if node.target == torch.ops.dc.allgather_param.default] + wait_nodes = [node for node in new_graph.nodes if node.target == torch.ops.dc.wait_allgather.default] + release_nodes = [node for node in new_graph.nodes if node.target == torch.ops.dc.release_param.default] + output_node = next(node for node in new_graph.nodes if node.op == "output") + + assert len(ag_nodes) == 1 + assert len(wait_nodes) == 1 + assert release_nodes == [] + assert ag_nodes[0].args[0].name == param.name + assert wait_nodes[0].args[0] is ag_nodes[0] + assert output_node.args == ((wait_nodes[0], ), ) + + +def test_zero3_scheduler_debug_logs_disabled_budget(monkeypatch, capsys): + monkeypatch.setenv(zero3_compile_mod.SCHEDULER_DEBUG_ENV, "1") + monkeypatch.setattr(zero3_compile_mod.dist, "is_initialized", lambda: False) + graph = Graph() + graph.output(()) + graph.lint() + + zero3_compile_mod._log_scheduler_result(7, + bwd=True, + scheduler_budget=None, + disabled_reason="missing_or_incomplete_memory_profile", + graph=graph) + + captured = capsys.readouterr() + assert "budget_enabled=False" in captured.out + assert "disabled_reason=missing_or_incomplete_memory_profile" in captured.out + + +def test_zero3_scheduler_debug_logs_smallest_rejected_candidate_peak(monkeypatch, capsys): + monkeypatch.setenv(zero3_compile_mod.SCHEDULER_DEBUG_ENV, "1") + monkeypatch.setattr(zero3_compile_mod.dist, "is_initialized", lambda: False) + graph = Graph() + graph.output(()) + graph.lint() + setattr( + graph, + schedule_mod.SCHEDULER_BUDGET_DIAGNOSTICS_ATTR, + { + "selected": [], + "budget_rejections": 2, + "budget_rejected_candidates": [{ + "over_budget_bytes": 64 + }, { + "over_budget_bytes": 32 + }], + "budget_overflows": [], + }, + ) + budget = schedule_mod.SchedulerMemoryBudget(max_gathered_bytes=1024, source="test") + + zero3_compile_mod._log_scheduler_result(7, bwd=True, scheduler_budget=budget, disabled_reason=None, graph=graph) + + captured = capsys.readouterr() + assert "minimum_rejected_candidate_peak_bytes=1056" in captured.out + + +def test_zero3_final_schedule_fingerprint_detects_rank_mismatch(monkeypatch): + monkeypatch.setenv(zero3_compile_mod.SCHEDULER_DEBUG_ENV, "1") + monkeypatch.setattr(zero3_compile_mod.dist, "is_initialized", lambda: True) + monkeypatch.setattr(zero3_compile_mod, "get_accelerator", lambda: SimpleNamespace(current_device=lambda: "cpu")) + messages = [] + monkeypatch.setattr( + zero3_compile_mod, + "_print_scheduler_debug", + lambda message, process_group=None: messages.append(message), + ) + + def reduce_mismatched_fingerprints(value, op): + value[0] = 1 if op == zero3_compile_mod.dist.ReduceOp.MIN else 2 + + monkeypatch.setattr(zero3_compile_mod.dist, "all_reduce", reduce_mismatched_fingerprints) + graph = Graph() + graph.output(()) + + with pytest.raises(RuntimeError, match="final schedule fingerprint mismatch"): + zero3_compile_mod._validate_final_schedule_fingerprint(graph, graph_id=7, bwd=False) + assert any("collective_schedule_projection" in message for message in messages) + + +def test_zero3_final_schedule_fingerprint_is_safe_without_distributed(monkeypatch, capsys): + monkeypatch.setenv(zero3_compile_mod.SCHEDULER_DEBUG_ENV, "1") + monkeypatch.setattr(zero3_compile_mod.dist, "is_initialized", lambda: False) + monkeypatch.setattr(zero3_compile_mod.dist, "all_reduce", + lambda *args, **kwargs: pytest.fail("non-distributed debug must not use a collective")) + graph = Graph() + graph.output(()) + + fingerprint = zero3_compile_mod._validate_final_schedule_fingerprint(graph, graph_id=8, bwd=True) + + assert fingerprint == zero3_compile_mod._final_schedule_fingerprint(graph) + output = capsys.readouterr().out + assert "collective_schedule_projection" in output + assert "final_schedule_fingerprint" in output + + +def test_zero3_collective_schedule_projection_ignores_names_and_graph_ids(): + + def make_graph(graph_id, name): + graph = Graph() + param = _placeholder(graph, f"{name}_param") + gather = _with_meta( + graph.call_function(torch.ops.dc.allgather_param.default, (param, graph_id, 17), name=f"{name}_gather")) + wait = _with_meta( + graph.call_function(torch.ops.dc.wait_allgather.default, (gather, graph_id, 17), name=f"{name}_wait")) + release = _with_meta( + graph.call_function(torch.ops.dc.release_param.default, (wait, graph_id, 17, 1), name=f"{name}_release")) + graph.output(release) + return graph + + first_projection = zero3_compile_mod._collective_schedule_projection(make_graph(1, "first")) + second_projection = zero3_compile_mod._collective_schedule_projection(make_graph(999, "second")) + + assert first_projection == [("gather", 17), ("release", 17)] + assert second_projection == first_projection + assert (zero3_compile_mod._collective_schedule_fingerprint(first_projection) == + zero3_compile_mod._collective_schedule_fingerprint(second_projection)) + + +def test_zero3_collective_schedule_projection_detects_order_divergence(): + graph = Graph() + first_param = _placeholder(graph, "first_param") + second_param = _placeholder(graph, "second_param") + first_gather = _allgather(graph, first_param, 17, "first") + second_gather = _allgather(graph, second_param, 23, "second") + first_release = _release(graph, first_gather, 17, "first") + second_release = _release(graph, second_gather, 23, "second") + graph.output((first_release, second_release)) + + projection = zero3_compile_mod._collective_schedule_projection(graph) + reordered = [projection[1], projection[0], *projection[2:]] + + assert projection[:2] == [("gather", 17), ("gather", 23)] + assert (zero3_compile_mod._collective_schedule_fingerprint(projection) + != zero3_compile_mod._collective_schedule_fingerprint(reordered)) + + +def test_zero3_final_schedule_fingerprint_is_absent_when_debug_is_disabled(monkeypatch): + monkeypatch.delenv(zero3_compile_mod.SCHEDULER_DEBUG_ENV, raising=False) + monkeypatch.delenv(zero3_compile_mod.SCHEDULER_DEBUG_ENV_LEGACY, raising=False) + monkeypatch.setattr(zero3_compile_mod, "_final_schedule_fingerprint", + lambda graph: pytest.fail("non-debug scheduling must not compute a fingerprint")) + graph = Graph() + graph.output(()) + + assert zero3_compile_mod._validate_final_schedule_fingerprint(graph, graph_id=9, bwd=False) is None + + +def test_zero3_scheduler_collectives_stay_with_each_data_parallel_group(monkeypatch): + monkeypatch.setenv(zero3_compile_mod.SCHEDULER_DEBUG_ENV, "1") + monkeypatch.setattr(zero3_compile_mod.dist, "is_initialized", lambda: True) + monkeypatch.setattr(zero3_compile_mod, "get_accelerator", + lambda: SimpleNamespace(current_device=lambda: "cpu", total_memory=lambda: 2000)) + group_a = object() + group_b = object() + group_sizes = {group_a: 2, group_b: 8} + collective_groups = [] + + def get_world_size(group=None): + assert group in group_sizes + return group_sizes[group] + + def all_reduce(tensor, op, group=None): + assert group in group_sizes + collective_groups.append(group) + + monkeypatch.setattr(zero3_compile_mod.dist, "get_world_size", get_world_size) + monkeypatch.setattr(zero3_compile_mod.dist, "get_rank", lambda group=None: 0) + monkeypatch.setattr(zero3_compile_mod.dist, "all_reduce", all_reduce) + monkeypatch.setattr(graph_profile_mod.dist, "all_reduce", all_reduce) + + def make_group_graph(name): + graph = Graph() + param = _placeholder(graph, f"{name}_param") + ag = _allgather(graph, param, 1, name, tensor_size=102) + wait = _wait(graph, ag, 1, name) + op = _neg(graph, wait, f"{name}_op") + release = _release(graph, op, 1, name) + graph.output((release, )) + graph.lint() + for node in graph.nodes: + node.meta.setdefault("alloc_mem", 0) + node.meta.setdefault("max_mem", 0) + node.meta.setdefault("profile_mem_start", 100) + node.meta.setdefault("profile_mem_peak", 1000) + return GraphModule(torch.nn.Module(), graph), ag + + gm_a, ag_a = make_group_graph("group_a") + gm_b, ag_b = make_group_graph("group_b_with_different_graph") + + zero3_compile_mod._scheduler_budget_from_operator_profile(gm_a, process_group=group_a) + zero3_compile_mod._validate_final_schedule_fingerprint(gm_a.graph, graph_id=10, bwd=False, process_group=group_a) + graph_profile_mod._rank_max_profile_memory(100, 200, torch.device("cpu"), distributed=True, process_group=group_a) + zero3_compile_mod._scheduler_budget_from_operator_profile(gm_b, process_group=group_b) + zero3_compile_mod._validate_final_schedule_fingerprint(gm_b.graph, graph_id=11, bwd=True, process_group=group_b) + graph_profile_mod._rank_max_profile_memory(300, 400, torch.device("cpu"), distributed=True, process_group=group_b) + + assert ag_a.meta["allgather_allocation_bytes"] == 104 + assert ag_b.meta["allgather_allocation_bytes"] == 112 + assert collective_groups + first_group_b = collective_groups.index(group_b) + assert all(group is group_a for group in collective_groups[:first_group_b]) + assert all(group is group_b for group in collective_groups[first_group_b:]) + + +def test_prefetch_and_selective_gather_collectives_stay_with_data_parallel_group(monkeypatch): + try: + torch.ops.dc.reload_parameter.default + except AttributeError: + library = torch.library.Library("dc", "FRAGMENT") + library.define("reload_parameter(Tensor a, int graph_id, int id) -> ()") + _DC_LIBRARIES.append(library) + + group_a = object() + group_b = object() + collective_groups = [] + + class FakeAccelerator: + + def current_device(self): + return "cpu" + + def total_memory(self): + return 2000 + + def available_memory(self): + return 1800 + + def memory_allocated(self): + return 0 + + def max_memory_allocated(self): + return 0 + + def all_reduce(tensor, op, group=None): + assert group in (group_a, group_b) + collective_groups.append(group) + + monkeypatch.setattr(prefetch_mod.dist, "all_reduce", all_reduce) + monkeypatch.setattr(prefetch_mod, "get_accelerator", FakeAccelerator) + monkeypatch.setattr(prefetch_mod, "create_predictor", lambda: lambda _: 0.0) + monkeypatch.setattr(prefetch_mod, "print_rank_0", lambda _: None) + monkeypatch.setattr(selective_gather_mod, "get_accelerator", FakeAccelerator) + monkeypatch.setattr(selective_gather_mod, "get_deepcompile_handle", + lambda: SimpleNamespace(set_persistent=lambda _: None)) + monkeypatch.setattr(selective_gather_mod, "print_rank_0", lambda _: None) + + def run_passes(group, graph_id): + graph = Graph() + value = _placeholder(graph, f"group_{graph_id}_input") + result = _neg(graph, value, f"group_{graph_id}_result") + graph.output((result, )) + graph.lint() + mem = [(node.name, 0, 0, 0) for node in graph.nodes] + timing = [(node.name, 0.0, 0.0) for node in graph.nodes] + tensor_sizes = [(node.name, 0) for node in graph.nodes] + profile = ProfilingResult(fwd_graph=graph, + bwd_graph=graph, + needs_backward=True, + fwd_mem=mem, + bwd_mem=mem, + fwd_time=timing, + bwd_time=timing, + fwd_tensor_sizes=tensor_sizes, + bwd_tensor_sizes=tensor_sizes, + process_group=group) + profiling_results = {graph_id: profile} + gm = GraphModule(torch.nn.Module(), graph) + prefetch_mod.schedule_prefetch(gm, + graph_id=graph_id, + graph_order=[(graph_id, True)], + profiling_results=profiling_results, + create_inputs_fn=lambda: (), + mem_budget=0, + param_manager={}, + bwd=False) + selective_gather_mod.selective_gather(gm, + graph_id=graph_id, + graph_order=[(graph_id, True)], + profiling_results=profiling_results, + create_inputs_fn=lambda: (), + mem_budget=0, + param_manager={}, + bwd=True) + + run_passes(group_a, 0) + run_passes(group_b, 1) + + assert collective_groups == [group_a, group_a, group_b, group_b] + + +def _allgather(graph, arg, ds_id, name, tensor_size=1, device_time=1, allocation_size=None): + node = _with_meta( graph.call_function(torch.ops.dc.allgather_param.default, (arg, 0, ds_id), {"dtype": torch.float16}, name=f"allgather_ds_param_{name}_{ds_id}"), tensor_size=tensor_size, device_time=device_time, ) + if allocation_size is not None: + node.meta["allgather_allocation_bytes"] = allocation_size + return node def _wait(graph, arg, ds_id, name): @@ -115,14 +968,28 @@ def _add(graph, lhs, rhs, name, device_time=0): return _with_meta(graph.call_function(operator.add, (lhs, rhs), name=name), device_time=device_time) -def _release(graph, arg, ds_id, name): +def _release(graph, arg, ds_id, name, n_users=1): return _with_meta( - graph.call_function(torch.ops.dc.release_param.default, (arg, 0, ds_id, 1), + graph.call_function(torch.ops.dc.release_param.default, (arg, 0, ds_id, n_users), name=f"release_ds_param_{name}_{ds_id}")) -def _scheduled_names(graph): - return [node.name for node in schedule_mod.fast_free_schedule(graph, 0, 0, debug_log=True).nodes] +def _scheduled_graph(graph, scheduler_budget=None): + return schedule_mod.fast_free_schedule( + graph, + 0, + 0, + debug_log=True, + scheduler_budget=scheduler_budget, + ) + + +def _scheduled_names(graph, scheduler_budget=None): + return [node.name for node in _scheduled_graph(graph, scheduler_budget=scheduler_budget).nodes] + + +def _scheduler_diagnostics(graph): + return getattr(graph, schedule_mod.SCHEDULER_BUDGET_DIAGNOSTICS_ATTR) def test_fast_free_schedule_keeps_zero_free_acc_filter(): @@ -217,6 +1084,292 @@ def test_fast_free_schedule_uses_pressure_tiebreaker_in_fallback_bucket(): assert names.index(low_ag.name) < names.index(high_ag.name) +def test_fast_free_schedule_zero_rejection_budget_preserves_legacy_order(): + graph = Graph() + + high_param = _placeholder(graph, "non_constraining_high_param") + high_extra_param = _placeholder(graph, "non_constraining_high_extra_param") + low_param = _placeholder(graph, "non_constraining_low_param") + low_extra_param = _placeholder(graph, "non_constraining_low_extra_param") + + high_ag = _allgather(graph, high_param, 70, "non_constraining_high", tensor_size=100, allocation_size=100) + high_wait = _wait(graph, high_ag, 70, "non_constraining_high") + high_extra_ag = _allgather(graph, + high_extra_param, + 71, + "non_constraining_high_extra", + tensor_size=10, + allocation_size=10) + high_extra_wait = _wait(graph, high_extra_ag, 71, "non_constraining_high_extra") + high_use = _add(graph, high_wait, high_extra_wait, "non_constraining_high_use", device_time=1) + high_release = _release(graph, high_use, 70, "non_constraining_high") + + low_ag = _allgather(graph, low_param, 80, "non_constraining_low", tensor_size=1, allocation_size=200) + low_wait = _wait(graph, low_ag, 80, "non_constraining_low") + low_extra_ag = _allgather(graph, + low_extra_param, + 81, + "non_constraining_low_extra", + tensor_size=10, + allocation_size=10) + low_extra_wait = _wait(graph, low_extra_ag, 81, "non_constraining_low_extra") + low_use = _add(graph, low_wait, low_extra_wait, "non_constraining_low_use", device_time=100) + low_release = _release(graph, low_use, 80, "non_constraining_low") + + graph.output((high_release, low_release)) + graph.lint() + + legacy_graph = _scheduled_graph(graph) + legacy_names = [node.name for node in legacy_graph.nodes] + budget = schedule_mod.SchedulerMemoryBudget(max_gathered_bytes=220, source="test") + assert budget.max_gathered_bytes < schedule_mod.max_possible_gathered_bytes(graph) + scheduled_graph = _scheduled_graph(graph, scheduler_budget=budget) + budget_names = [node.name for node in scheduled_graph.nodes] + + assert _scheduler_diagnostics(scheduled_graph)["budget_rejections"] == 0 + assert budget_names == legacy_names + assert (zero3_compile_mod._final_schedule_fingerprint(scheduled_graph) == + zero3_compile_mod._final_schedule_fingerprint(legacy_graph)) + + +def test_fast_free_schedule_counts_live_gathered_bytes_when_filtering_candidates(): + graph = Graph() + + first_param = _placeholder(graph, "budget_first_param") + high_param = _placeholder(graph, "budget_high_param") + low_param = _placeholder(graph, "budget_low_param") + + first_ag = _allgather(graph, first_param, 80, "budget_first", tensor_size=70) + first_wait = _wait(graph, first_ag, 80, "budget_first") + + high_dep = _add(graph, high_param, first_wait, "budget_high_dep") + high_ag = _allgather(graph, high_dep, 81, "budget_high", tensor_size=40) + high_wait = _wait(graph, high_ag, 81, "budget_high") + high_use = _neg(graph, high_wait, "budget_high_use", device_time=1) + high_release = _release(graph, high_use, 81, "budget_high") + + low_dep = _add(graph, low_param, first_wait, "budget_low_dep") + low_ag = _allgather(graph, low_dep, 82, "budget_low", tensor_size=20) + low_wait = _wait(graph, low_ag, 82, "budget_low") + low_use = _neg(graph, low_wait, "budget_low_use", device_time=100) + low_release = _release(graph, low_use, 82, "budget_low") + + high_low_pair = _add(graph, high_wait, low_wait, "budget_high_low_pair") + first_use = _add(graph, first_wait, high_low_pair, "budget_first_use") + first_release = _release(graph, first_use, 80, "budget_first") + + graph.output((first_release, high_release, low_release)) + graph.lint() + + no_budget_names = _scheduled_names(graph) + assert no_budget_names.index(high_ag.name) < no_budget_names.index(low_ag.name) + + budget = schedule_mod.SchedulerMemoryBudget(max_gathered_bytes=100, source="test") + scheduled_graph = _scheduled_graph(graph, scheduler_budget=budget) + names = [node.name for node in scheduled_graph.nodes] + diagnostics = _scheduler_diagnostics(scheduled_graph) + + assert names.index(first_ag.name) < names.index(low_ag.name) + assert names.index(low_ag.name) < names.index(high_ag.name) + assert diagnostics["budget_rejections"] > 0 + assert any(record["node"] == high_ag.name for record in diagnostics["budget_rejected_candidates"]) + + +def test_fast_free_schedule_continues_to_higher_count_candidates_when_lowest_count_exceeds_budget(): + graph = Graph() + + first_param = _placeholder(graph, "budget_count_first_param") + high_param = _placeholder(graph, "budget_count_high_param") + low_param = _placeholder(graph, "budget_count_low_param") + extra_param = _placeholder(graph, "budget_count_extra_param") + + first_ag = _allgather(graph, first_param, 100, "budget_count_first", tensor_size=70) + first_wait = _wait(graph, first_ag, 100, "budget_count_first") + + high_dep = _add(graph, high_param, first_wait, "budget_count_high_dep") + high_ag = _allgather(graph, high_dep, 101, "budget_count_high", tensor_size=60) + high_wait = _wait(graph, high_ag, 101, "budget_count_high") + high_use = _neg(graph, high_wait, "budget_count_high_use") + high_release = _release(graph, high_use, 101, "budget_count_high") + + low_dep = _add(graph, low_param, first_wait, "budget_count_low_dep") + low_ag = _allgather(graph, low_dep, 102, "budget_count_low", tensor_size=20) + low_wait = _wait(graph, low_ag, 102, "budget_count_low") + extra_dep = _add(graph, extra_param, low_wait, "budget_count_extra_dep") + extra_ag = _allgather(graph, extra_dep, 103, "budget_count_extra", tensor_size=20) + extra_wait = _wait(graph, extra_ag, 103, "budget_count_extra") + low_use = _add(graph, low_wait, extra_wait, "budget_count_low_use") + low_release = _release(graph, low_use, 102, "budget_count_low") + extra_release = _release(graph, low_use, 103, "budget_count_extra") + + first_use = _add(graph, first_wait, low_wait, "budget_count_first_use") + first_release = _release(graph, first_use, 100, "budget_count_first") + + graph.output((first_release, high_release, low_release, extra_release)) + graph.lint() + + budget = schedule_mod.SchedulerMemoryBudget(max_gathered_bytes=120, source="test") + names = _scheduled_names(graph, scheduler_budget=budget) + + assert names.index(low_ag.name) < names.index(high_ag.name) + + +def test_fast_free_schedule_counts_padded_allgather_allocation_bytes(): + graph = Graph() + + param = _placeholder(graph, "budget_padded_param") + ag = _allgather(graph, param, 110, "budget_padded", tensor_size=102, allocation_size=112) + wait = _wait(graph, ag, 110, "budget_padded") + use = _neg(graph, wait, "budget_padded_use") + release = _release(graph, use, 110, "budget_padded") + + graph.output((release, )) + graph.lint() + + budget = schedule_mod.SchedulerMemoryBudget(max_gathered_bytes=105, source="test") + scheduled_graph = _scheduled_graph(graph, scheduler_budget=budget) + diagnostics = _scheduler_diagnostics(scheduled_graph) + + assert diagnostics["budget_overflows"][0]["candidate_allgather_bytes"] == 112 + assert diagnostics["budget_overflows"][0]["over_budget_bytes"] == 7 + assert diagnostics["budget_overflows"][0]["path"] == "until_free" + + +def test_fast_free_schedule_records_diagnostic_when_no_candidate_fits_budget(): + graph = Graph() + + first_param = _placeholder(graph, "budget_fail_first_param") + high_param = _placeholder(graph, "budget_fail_high_param") + low_param = _placeholder(graph, "budget_fail_low_param") + + first_ag = _allgather(graph, first_param, 90, "budget_fail_first", tensor_size=80) + first_wait = _wait(graph, first_ag, 90, "budget_fail_first") + + high_dep = _add(graph, high_param, first_wait, "budget_fail_high_dep") + high_ag = _allgather(graph, high_dep, 91, "budget_fail_high", tensor_size=40) + high_wait = _wait(graph, high_ag, 91, "budget_fail_high") + high_use = _neg(graph, high_wait, "budget_fail_high_use", device_time=1) + high_release = _release(graph, high_use, 91, "budget_fail_high") + + low_dep = _add(graph, low_param, first_wait, "budget_fail_low_dep") + low_ag = _allgather(graph, low_dep, 92, "budget_fail_low", tensor_size=30) + low_wait = _wait(graph, low_ag, 92, "budget_fail_low") + low_use = _neg(graph, low_wait, "budget_fail_low_use", device_time=100) + low_release = _release(graph, low_use, 92, "budget_fail_low", n_users=2) + + first_use = _add(graph, first_wait, low_wait, "budget_fail_first_use") + first_release = _release(graph, first_use, 90, "budget_fail_first") + low_last_release = _release(graph, first_use, 92, "budget_fail_low_last", n_users=2) + + graph.output((first_release, high_release, low_release, low_last_release)) + graph.lint() + + budget = schedule_mod.SchedulerMemoryBudget(max_gathered_bytes=100, source="test") + scheduled_graph = _scheduled_graph(graph, scheduler_budget=budget) + names = [node.name for node in scheduled_graph.nodes] + diagnostics = _scheduler_diagnostics(scheduled_graph) + + assert names.index(first_ag.name) < names.index(low_ag.name) + assert diagnostics["budget_rejections"] > 0 + assert diagnostics["budget_overflows"][0]["source"] == "test" + assert diagnostics["budget_overflows"][0]["max_gathered_bytes"] == 100 + assert diagnostics["budget_overflows"][0]["over_budget_bytes"] > 0 + + +def test_fast_free_schedule_over_budget_fallback_prefers_lower_peak_before_live_memory(): + graph = Graph() + + first_param = _placeholder(graph, "budget_debt_first_param") + helper_param = _placeholder(graph, "budget_debt_helper_param") + + first_ag = _allgather(graph, first_param, 120, "budget_debt_first", tensor_size=80) + first_wait = _wait(graph, first_ag, 120, "budget_debt_first") + helper_dep = _add(graph, helper_param, first_wait, "budget_debt_helper_dep") + helper_ag = _allgather(graph, helper_dep, 121, "budget_debt_helper", tensor_size=30) + helper_wait = _wait(graph, helper_ag, 121, "budget_debt_helper") + use = _add(graph, first_wait, helper_wait, "budget_debt_use") + first_release = _release(graph, use, 120, "budget_debt_first") + helper_release = _release(graph, use, 121, "budget_debt_helper") + + graph.output((first_release, helper_release)) + graph.lint() + + budget = schedule_mod.SchedulerMemoryBudget(max_gathered_bytes=50, source="test") + scheduled_graph = _scheduled_graph(graph, scheduler_budget=budget) + diagnostics = _scheduler_diagnostics(scheduled_graph) + first_selection = diagnostics["selected"][0] + + assert first_selection["path"] == "until_ag" + assert first_selection["schedule_until_ag_peak_mem"] < first_selection["schedule_until_free_peak_mem"] + assert first_selection["schedule_until_ag_live_mem"] > first_selection["schedule_until_free_live_mem"] + assert diagnostics["budget_overflows"][0]["path"] == "until_ag" + + +def test_over_budget_fallback_prefers_lower_peak_before_ending_residency(): + graph = Graph() + high_peak_node = _placeholder(graph, "high_peak_zero_residency") + low_peak_node = _placeholder(graph, "low_peak_nonzero_residency") + + high_peak_task = schedule_mod.AllgatherTask(node=high_peak_node, + allgather_cost=0, + free_cost=0, + allgathered_mem=1000, + allgather_acc_mem=1000, + free_acc_mem=0, + last_use=high_peak_node, + n_scheduled_ags=1, + schedule_until_ag=[high_peak_node], + schedule_until_free=[high_peak_node], + schedule_until_ag_peak_mem=1000, + schedule_until_free_peak_mem=1000, + schedule_until_ag_live_mem=1000, + schedule_until_free_live_mem=0) + low_peak_task = schedule_mod.AllgatherTask(node=low_peak_node, + allgather_cost=0, + free_cost=0, + allgathered_mem=51, + allgather_acc_mem=51, + free_acc_mem=0, + last_use=low_peak_node, + n_scheduled_ags=1, + schedule_until_ag=[low_peak_node], + schedule_until_free=[low_peak_node], + schedule_until_ag_peak_mem=51, + schedule_until_free_peak_mem=51, + schedule_until_ag_live_mem=51, + schedule_until_free_live_mem=51) + + selected, _ = schedule_mod._select_over_budget_allgather_task([high_peak_task, low_peak_task], + schedule_mod.SchedulerMemoryBudget( + max_gathered_bytes=50, source="test")) + + assert selected is low_peak_task + + +def test_candidate_peak_resets_after_overflow_is_released(): + graph = Graph() + first_param = _placeholder(graph, "historical_overflow_param") + next_param = _placeholder(graph, "later_fitting_param") + first_ag = _allgather(graph, first_param, 130, "historical_overflow", allocation_size=100) + first_release = _release(graph, first_ag, 130, "historical_overflow") + next_ag = _allgather(graph, next_param, 131, "later_fitting", allocation_size=40) + + tracker = schedule_mod._GatheredParamTracker({130: 1, 131: 1}) + tracker.apply(first_ag) + tracker.apply(first_release) + assert tracker.live_bytes == 0 + assert tracker.peak_bytes == 100 + + candidate_peak, candidate_live = schedule_mod._simulate_path_stats(tracker, [next_ag]) + + assert candidate_peak == 40 + assert candidate_live == 40 + assert schedule_mod._fits_budget(schedule_mod.SchedulerMemoryBudget(max_gathered_bytes=50, source="test"), + candidate_peak) + assert tracker.live_bytes == 0 + assert tracker.peak_bytes == 100 + + def test_fast_free_schedule_keeps_single_allgather_release_order(): graph = Graph() diff --git a/tests/unit/compile/test_util.py b/tests/unit/compile/test_util.py new file mode 100644 index 000000000000..0422fefb7576 --- /dev/null +++ b/tests/unit/compile/test_util.py @@ -0,0 +1,71 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import deepspeed.compile.util as compile_util + + +def test_all_reduce_preserves_default_and_explicit_group_forwarding(monkeypatch): + calls = [] + sentinel = object() + + def unexpected_is_initialized(): + raise AssertionError("all_reduce must not consult distributed initialization state") + + def record_all_reduce(*args, **kwargs): + calls.append((args, kwargs)) + return sentinel + + monkeypatch.setattr(compile_util.dist, "is_initialized", unexpected_is_initialized) + monkeypatch.setattr(compile_util.dist, "all_reduce", record_all_reduce) + + tensor = object() + op = object() + process_group = object() + + assert compile_util.all_reduce(tensor, op) is sentinel + assert compile_util.all_reduce(tensor, op, process_group) is sentinel + assert calls == [((tensor, op), {}), ((tensor, op), {"group": process_group})] + + +def test_get_rank_preserves_default_and_explicit_group_forwarding(monkeypatch): + calls = [] + sentinel = object() + + def unexpected_is_initialized(): + raise AssertionError("get_rank must not consult distributed initialization state") + + def record_get_rank(*args, **kwargs): + calls.append((args, kwargs)) + return sentinel + + monkeypatch.setattr(compile_util.dist, "is_initialized", unexpected_is_initialized) + monkeypatch.setattr(compile_util.dist, "get_rank", record_get_rank) + + process_group = object() + + assert compile_util.get_rank() is sentinel + assert compile_util.get_rank(process_group) is sentinel + assert calls == [((), {}), ((), {"group": process_group})] + + +def test_barrier_preserves_default_and_explicit_group_forwarding(monkeypatch): + calls = [] + sentinel = object() + + def unexpected_is_initialized(): + raise AssertionError("barrier must not consult distributed initialization state") + + def record_barrier(*args, **kwargs): + calls.append((args, kwargs)) + return sentinel + + monkeypatch.setattr(compile_util.dist, "is_initialized", unexpected_is_initialized) + monkeypatch.setattr(compile_util.dist, "barrier", record_barrier) + + process_group = object() + + assert compile_util.barrier() is sentinel + assert compile_util.barrier(process_group) is sentinel + assert calls == [((), {}), ((), {"group": process_group})] diff --git a/tests/unit/compile/test_z3_eager_fallback.py b/tests/unit/compile/test_z3_eager_fallback.py new file mode 100644 index 000000000000..e9ceb2d7b406 --- /dev/null +++ b/tests/unit/compile/test_z3_eager_fallback.py @@ -0,0 +1,76 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from types import FunctionType + +import torch + +from deepspeed.compile import z3_eager_fallback +from deepspeed.compile.z3_eager_fallback import DeepCompileZ3EagerFallback +from deepspeed.runtime.zero import parameter_offload +from deepspeed.runtime.zero.parameter_offload import ZeROOrderedDict +from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus + + +def _zero_module_with_param(): + module = torch.nn.Module() + param = torch.nn.Parameter(torch.empty(0)) + param.ds_id = 7 + param.ds_status = ZeroParamStatus.NOT_AVAILABLE + + params = ZeROOrderedDict(parent_module=module) + params["weight"] = param + params._in_forward = True + module._parameters = params + return module, param + + +def test_deepcompile_fallback_suppresses_guard_time_gather(monkeypatch): + module, param = _zero_module_with_param() + fallback = DeepCompileZ3EagerFallback(engine=None) + gather_calls = [] + + param.all_gather = lambda: gather_calls.append(param.ds_id) + monkeypatch.setattr(z3_eager_fallback, "_ACTIVE_FALLBACK", fallback) + monkeypatch.setattr(z3_eager_fallback, "is_dynamo_guard_evaluation", lambda: True) + + assert module._parameters["weight"] is param + assert gather_calls == [] + assert param.ds_status == ZeroParamStatus.NOT_AVAILABLE + + +def test_deepcompile_fallback_still_gathers_outside_guard(monkeypatch): + module, param = _zero_module_with_param() + fallback = DeepCompileZ3EagerFallback(engine=None) + gather_calls = [] + + def all_gather(): + gather_calls.append(param.ds_id) + param.ds_status = ZeroParamStatus.AVAILABLE + + param.all_gather = all_gather + monkeypatch.setattr(z3_eager_fallback, "_ACTIVE_FALLBACK", fallback) + monkeypatch.setattr(z3_eager_fallback, "is_dynamo_guard_evaluation", lambda: False) + monkeypatch.setattr(parameter_offload, "print_rank_0", lambda *args, **kwargs: None) + + assert module._parameters["weight"] is param + assert gather_calls == [7] + assert fallback.stats()["last_gathered_param_ids"] == [7] + + +def test_dynamo_guard_detection_is_false_outside_guard_stack(): + assert not z3_eager_fallback.is_dynamo_guard_evaluation() + + +def test_dynamo_guard_detection_is_true_inside_guard_module_stack(): + guard_fn = FunctionType( + compile("result = is_dynamo_guard_evaluation()", "", "exec"), + { + "__name__": "torch._dynamo.guards", + "is_dynamo_guard_evaluation": z3_eager_fallback.is_dynamo_guard_evaluation, + }, + ) + guard_fn() + assert guard_fn.__globals__["result"] is True diff --git a/tests/unit/compile/test_zero3_grad_dtype.py b/tests/unit/compile/test_zero3_grad_dtype.py index a2fdaae89d9d..37ae8152b7ca 100644 --- a/tests/unit/compile/test_zero3_grad_dtype.py +++ b/tests/unit/compile/test_zero3_grad_dtype.py @@ -3,9 +3,12 @@ # DeepSpeed Team +import pytest import torch -from deepspeed.compile.init_z3 import _resolve_expected_grad_dtype +from deepspeed.compile import backend as backend_mod +from deepspeed.compile.init_z3 import _allow_dynamo_dynamic_parameter_shapes_for_z3, _resolve_expected_grad_dtype +from deepspeed.runtime.engine import DeepSpeedEngine def test_missing_grad_dtype_attribute_falls_back_to_param_dtype(): @@ -28,3 +31,183 @@ def test_explicit_grad_dtype_is_preserved(): param.grad_dtype = torch.float32 assert _resolve_expected_grad_dtype(param) is torch.float32 + + +def test_zero3_allows_dynamo_dynamic_parameter_shapes(monkeypatch): + + class FakeDynamoConfig: + force_parameter_static_shapes = True + force_nn_module_property_static_shapes = True + + class FakeDynamo: + config = FakeDynamoConfig() + + monkeypatch.setattr(torch, "_dynamo", FakeDynamo) + + restore = _allow_dynamo_dynamic_parameter_shapes_for_z3({}) + assert restore + try: + assert FakeDynamo.config.force_parameter_static_shapes is False + assert FakeDynamo.config.force_nn_module_property_static_shapes is False + finally: + restore() + + +def test_zero3_preserves_explicit_dynamo_dynamic_setting(monkeypatch): + + class FakeDynamoConfig: + force_parameter_static_shapes = True + force_nn_module_property_static_shapes = True + + class FakeDynamo: + config = FakeDynamoConfig() + + compile_kwargs = {"dynamic": False} + monkeypatch.setattr(torch, "_dynamo", FakeDynamo) + + restore = _allow_dynamo_dynamic_parameter_shapes_for_z3(compile_kwargs) + assert restore + try: + assert compile_kwargs["dynamic"] is False + finally: + restore() + + +@pytest.mark.parametrize("first_owner_to_restore", [0, 1]) +def test_zero3_dynamo_config_restores_after_last_overlapping_owner(monkeypatch, first_owner_to_restore): + + class FakeDynamoConfig: + force_parameter_static_shapes = True + force_nn_module_property_static_shapes = False + + class FakeDynamo: + config = FakeDynamoConfig() + + monkeypatch.setattr(torch, "_dynamo", FakeDynamo) + restores = [_allow_dynamo_dynamic_parameter_shapes_for_z3({}), _allow_dynamo_dynamic_parameter_shapes_for_z3({})] + + assert all(restores) + restores[first_owner_to_restore]() + assert FakeDynamo.config.force_parameter_static_shapes is False + restores[1 - first_owner_to_restore]() + assert FakeDynamo.config.force_parameter_static_shapes is True + assert FakeDynamo.config.force_nn_module_property_static_shapes is False + + +@pytest.mark.parametrize("first_owner_to_destroy", [0, 1]) +def test_zero3_dynamo_config_restores_when_overlapping_engines_are_destroyed(monkeypatch, first_owner_to_destroy): + + class FakeDynamoConfig: + force_parameter_static_shapes = True + force_nn_module_property_static_shapes = False + + class FakeDynamo: + config = FakeDynamoConfig() + + monkeypatch.setattr(torch, "_dynamo", FakeDynamo) + engines = [object.__new__(DeepSpeedEngine), object.__new__(DeepSpeedEngine)] + for engine in engines: + torch.nn.Module.__init__(engine) + engine._deepcompile_active = False + engine._deepcompile_dynamo_config_restore = _allow_dynamo_dynamic_parameter_shapes_for_z3({}) + + engines[first_owner_to_destroy].destroy() + assert FakeDynamo.config.force_parameter_static_shapes is False + engines[1 - first_owner_to_destroy].destroy() + assert FakeDynamo.config.force_parameter_static_shapes is True + assert FakeDynamo.config.force_nn_module_property_static_shapes is False + + +@pytest.mark.parametrize("first_owner_to_destroy", [0, 1]) +def test_destroy_releases_only_owner_with_overlapping_frame_ids(first_owner_to_destroy): + original_autograd_function = torch.autograd.Function + engines = [object.__new__(DeepSpeedEngine), object.__new__(DeepSpeedEngine)] + owners = [object(), object()] + frame_id = 17 + for owner, engine in zip(owners, engines): + torch.nn.Module.__init__(engine) + engine._deepcompile_active = False + engine._deepcompile_owned_frames = {(owner, frame_id)} + + backend_mod.frames_needing_bwd.clear() + backend_mod.frames_needing_bwd.update(((owners[0], frame_id), (owners[1], frame_id))) + backend_mod.patch_compiled_func() + backend_mod.get_backward_inputs().append((torch.ones(1), )) + + try: + engines[first_owner_to_destroy].destroy() + surviving_owner = owners[1 - first_owner_to_destroy] + assert backend_mod.frames_needing_bwd == {(surviving_owner, frame_id)} + assert len(backend_mod.get_backward_inputs()) == 1 + assert torch.autograd.Function is not original_autograd_function + + engines[1 - first_owner_to_destroy].destroy() + assert backend_mod.frames_needing_bwd == set() + assert backend_mod.get_backward_inputs() == [] + assert torch.autograd.Function is original_autograd_function + finally: + backend_mod.frames_needing_bwd.clear() + backend_mod.unpatch_compiled_func() + + +def test_deactivation_releases_only_the_engine_owned_state(monkeypatch): + + class FakeDynamoConfig: + force_parameter_static_shapes = True + force_nn_module_property_static_shapes = False + + class FakeDynamo: + config = FakeDynamoConfig() + + monkeypatch.setattr(torch, "_dynamo", FakeDynamo) + engine = object.__new__(DeepSpeedEngine) + torch.nn.Module.__init__(engine) + engine._deepcompile_active = True + engine.module_forward_pre_hook = object() + engine.module_forward_post_hook = object() + engine._deepcompile_dynamo_config_restore = _allow_dynamo_dynamic_parameter_shapes_for_z3({}) + original_autograd_function = torch.autograd.Function + owner = object() + other_owner = object() + backend_mod.frames_needing_bwd.clear() + backend_mod.frames_needing_bwd.update(((owner, 17), (other_owner, 18))) + engine._deepcompile_owned_frames = {(owner, 17)} + backend_mod.patch_compiled_func() + backend_mod.get_backward_inputs().append((torch.ones(1), )) + + try: + engine._set_deepcompile_active(False) + + assert FakeDynamo.config.force_parameter_static_shapes is True + assert FakeDynamo.config.force_nn_module_property_static_shapes is False + assert not hasattr(engine, "_deepcompile_dynamo_config_restore") + assert engine._deepcompile_owned_frames == set() + assert backend_mod.frames_needing_bwd == {(other_owner, 18)} + assert len(backend_mod.get_backward_inputs()) == 1 + assert torch.autograd.Function is not original_autograd_function + assert engine.is_deepcompile_active() is False + finally: + backend_mod.frames_needing_bwd.clear() + backend_mod.unpatch_compiled_func() + + +def test_repeated_engine_destroy_cleans_native_state_once_and_deactivates(monkeypatch): + engine = object.__new__(DeepSpeedEngine) + torch.nn.Module.__init__(engine) + engine._deepcompile_active = True + cleanup_calls = [] + engine._release_deepcompile_compiled_backward_state = lambda: None + engine._release_deepcompile_dynamo_config = lambda: None + engine.is_deepcompile_active = lambda: engine._deepcompile_active + engine._set_deepcompile_active = lambda active: setattr(engine, "_deepcompile_active", active) + + fake_handle = type("FakeDeepCompileHandle", (), {"cleanup": lambda self: cleanup_calls.append("cleanup")})() + monkeypatch.setattr("deepspeed.runtime.engine.get_deepcompile_handle", lambda: fake_handle) + monkeypatch.setattr("deepspeed.runtime.engine.debug_clear_module_and_param_names", lambda: None) + + engine.destroy() + engine.destroy() + engine.__del__() + + assert cleanup_calls == ["cleanup"] + assert engine._deepcompile_active is False diff --git a/tests/unit/utils/test_allocator_telemetry.py b/tests/unit/utils/test_allocator_telemetry.py new file mode 100644 index 000000000000..c146b08ec5f6 --- /dev/null +++ b/tests/unit/utils/test_allocator_telemetry.py @@ -0,0 +1,72 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +import json + +from deepspeed.utils import allocator_telemetry as telemetry + + +def test_disabled_empty_cache_is_a_transparent_call(monkeypatch): + monkeypatch.delenv("DEEPSPEED_ALLOCATOR_TELEMETRY", raising=False) + telemetry._reset_for_tests() + calls = [] + + assert telemetry.record_empty_cache("test.site", lambda value: calls.append(value) or "ok", 7) == "ok" + assert calls == [7] + assert telemetry.summary()["empty_cache_counts"] == {} + + +def test_enabled_empty_cache_records_site_and_memory(monkeypatch): + monkeypatch.setenv("DEEPSPEED_ALLOCATOR_TELEMETRY", "1") + telemetry._reset_for_tests() + samples = iter([{"num_alloc_retries": 2}, {"num_alloc_retries": 3}, {"num_alloc_retries": 3}]) + monkeypatch.setattr(telemetry, "memory_stats", lambda: next(samples)) + telemetry.set_step(11, "measured") + + telemetry.record_empty_cache("backend.post-profile", lambda: None) + result = telemetry.summary() + + assert result["empty_cache_counts"] == {"backend.post-profile": 1} + assert result["empty_cache_observations"] == [{ + "site": "backend.post-profile", + "ordinal": 1, + "step": 11, + "phase": "measured", + "before": { + "num_alloc_retries": 2 + }, + "after": { + "num_alloc_retries": 3 + }, + "error": None, + }] + + +def test_scheduler_and_allocator_retry_terminal_summary(monkeypatch, capsys): + monkeypatch.setenv("DEEPSPEED_ALLOCATOR_TELEMETRY", "1") + telemetry._reset_for_tests() + monkeypatch.setattr(telemetry, "memory_stats", lambda: {"num_alloc_retries": 4}) + telemetry.set_step(12, "measured") + telemetry.record_scheduler_decision(graph_id=3, + backward=False, + budget_source="incomplete_profile_minimum_gather_residency", + disabled_reason=None, + max_gathered_bytes=1024, + max_live_gathered_bytes=1024, + budget_rejections=2, + over_budget_fallbacks=0, + minimum_rejected_candidate_peak_bytes=1536) + telemetry.record_allocator_retry_sample(3, 4) + + telemetry.close_recorder() + output = capsys.readouterr().out.strip() + prefix = "DEEPSPEED_ALLOCATOR_TELEMETRY_SUMMARY " + assert output.startswith(prefix) + payload = json.loads(output[len(prefix):]) + + assert payload["allocator_retry_total"] == 1 + assert payload["scheduler_decisions"][0]["budget_source"] == ("incomplete_profile_minimum_gather_residency") + assert payload["scheduler_decisions"][0]["minimum_rejected_candidate_peak_bytes"] == 1536 + assert payload["final_memory"]["num_alloc_retries"] == 4 diff --git a/tests/unit/v1/compile/test_graph_profile.py b/tests/unit/v1/compile/test_graph_profile.py index 3257b5bff980..1d5f1447ee7e 100644 --- a/tests/unit/v1/compile/test_graph_profile.py +++ b/tests/unit/v1/compile/test_graph_profile.py @@ -167,6 +167,7 @@ def raise_from_run(self, *args): assert interpreter.run() is None assert not interpreter.profile_complete assert interpreter.mem_record == [] + assert graph_profile.is_profile_incomplete(interpreter.graph) assert fake_handle.events == [("enable", True), ("clear", None), ("enable", False)] @@ -186,8 +187,36 @@ def fail_clear(): monkeypatch.setattr(graph_profile.Interpreter, "run", lambda self, *args: None) interpreter = graph_profile.MemoryProfilingInterpreter(_make_empty_graph_module()) + interpreter.env["retained"] = object() with pytest.raises(RuntimeError, match="cleanup failed"): interpreter.run() assert fake_handle.events == [("enable", True), ("clear", None), ("enable", False)] + assert interpreter.env == {} + + +def test_profiling_interpreter_restores_state_if_gathered_param_cleanup_fails(monkeypatch): + fake_handle = FakeDeepCompileHandle() + + def fail_clear(): + fake_handle.events.append(("clear", None)) + raise RuntimeError("cleanup failed") + + fake_handle.clear_all_gathered_params = fail_clear + + monkeypatch.setattr(graph_profile, "get_deepcompile_handle", lambda: fake_handle) + monkeypatch.setattr(graph_profile, "get_accelerator", lambda: FakeAccelerator()) + monkeypatch.setattr(graph_profile, "_all_real_if_tensor", lambda args: True) + monkeypatch.setattr(graph_profile, "_get_mem_usage_out_of_torch", lambda: 0) + monkeypatch.setattr(graph_profile.dist, "is_initialized", lambda: False) + monkeypatch.setattr(graph_profile.Interpreter, "run", lambda self, *args: None) + + interpreter = graph_profile.ProfilingInterpreter(_make_empty_graph_module(), iteration=1) + interpreter.env["retained"] = object() + + with pytest.raises(RuntimeError, match="cleanup failed"): + interpreter.run() + + assert fake_handle.events == [("enable", True), ("clear", None), ("enable", False)] + assert interpreter.env == {} diff --git a/tests/unit/v1/compile/test_selective_gather.py b/tests/unit/v1/compile/test_selective_gather.py index 02295f020872..3f498b2bc49a 100644 --- a/tests/unit/v1/compile/test_selective_gather.py +++ b/tests/unit/v1/compile/test_selective_gather.py @@ -152,7 +152,7 @@ def test_selective_gather_uses_profiled_headroom_instead_of_current_available_me assert fake_handle.persistent_ds_ids == [1, 2] -def test_selective_gather_uses_profiled_headroom_when_current_available_memory_is_low(monkeypatch): +def test_selective_gather_caps_profiled_headroom_by_current_available_memory(monkeypatch): fake_handle = FakeDeepCompileHandle() monkeypatch.setattr(selective_gather_pass, "get_accelerator", lambda: FakeAccelerator(available_mem=100)) @@ -192,7 +192,7 @@ def test_selective_gather_uses_profiled_headroom_when_current_available_memory_i bwd=True) assert returned is gm - assert fake_handle.persistent_ds_ids == [1, 2] + assert fake_handle.persistent_ds_ids == [] def test_selective_gather_skips_persistence_when_memory_profile_incomplete(monkeypatch):