Repository navigation
llm_shared/optimization: two implementations of meta-device layer eviction, one hand-rolled #13404
Description
Activity
Not consolidating this blind — the two implementations are not equivalent
I picked this up and stopped short of the swap, deliberately.
The consolidation looks like a one-liner: have
layer_inference._move_to_metadelegate tometa_eviction.evict_layer_to_meta. But the two differ in a way that matters for the thing eviction exists to do._evict_standard_layercallslayer.to("meta"), which goes throughnn.Module._applyand preserves parameter sharing — two names pointing at one tensor still point at one tensor afterwards._move_to_metawalksnamed_parameters(remove_duplicate=False)and re-registers every alias separately, deliberately untying shared storage. Its comment says why:a missed alias would keep the real tensor — and its memory — alive, defeating the eviction
So the hand-rolled version is not simply a worse copy of the canonical one; it encodes a different decision about aliased weights. Whether
.to("meta")frees the same memory for a layer with tied weights is a question about PyTorch's storage handling that I cannot answer from reading, and cannot test here: neithertorchnoraccelerateis installed on this box, andlayer_inference_test.pyis CI-only for that reason.Shipping an unverifiable behavioural change to a memory-eviction path, where the failure mode is a silent leak rather than an exception, is worse than leaving a documented duplication. So this stays open with the analysis recorded rather than half-done.
What would make it safe
- Determine whether
layer.to("meta")releases storage for tied weights, or whether the alias must be untied first — a short experiment on a machine with torch, using a module with deliberately shared parameters. - If
.to("meta")is sufficient: delete_move_to_meta,_set_parameterand_set_buffer, and callevict_layer_to_meta. Keep the tied-weight case as a test. - If it is not: keep both, and replace this issue's premise with a comment in
meta_eviction.pyexplaining why the hand-rolled path exists — so the next person does not "simplify" it either.
Either outcome ends the duplication. Only one of them can be chosen with evidence, and that evidence needs torch.
Related: #13495 fixed the bug that was living in the hand-rolled copy (
GrowingKVCachedtype coercion was a sibling of this), and PR #13402 first surfaced the divergence.- Determine whether
Summary
There are two independent implementations of "move a transformer layer to the meta device" in
autobot-backend/llm_shared/optimization/. One delegates to the supported PyTorch API; the other hand-rolls it, and that hand-rolled copy carried a bug that made layer eviction raise instead of free memory.Found while fixing
layer_inference_test.pyfor #13162 (PR #13402).The two implementations
Canonical —
meta_eviction.py(#1952):plus an
accelerate.set_module_tensor_to_devicepath for quantized modules. This is the supported API and handles parameters, buffers and shared storage correctly.Hand-rolled —
layer_inference.py::_move_to_meta(#1946):Walks
named_parameters()andnamed_buffers()itself and rebuilds each tensor on the meta device. This reimplements whatnn.Module._applyalready does.Why the duplication mattered
The hand-rolled copy did:
PyTorch rejects a cross-device-type
.dataassignment withRuntimeError: Attempted to call variable.set_data(tensor), but variable and tensor have incompatible tensor type, because a meta tensor carries a differentTensorImpl.LayerInferenceEngine.evict_layer()therefore raised on every call under real PyTorch — the eviction that the whole layer-by-layer mode exists to perform.nn.Module.to("meta")gets this right precisely becauseModule._applychecks_has_compatible_shallow_copy_typefirst and re-registers a newParameterwhen the types differ, instead of mutating.data.PR #13402 fixes the crash in place (re-registering parameters, mirroring the
_set_bufferhelper the same function already used), because that is the surgical change and the torch-dependent tests can only be verified in CI. It does not consolidate the two implementations.Proposed work
layer_inference._move_to_metaontometa_eviction.evict_layer_to_meta, or ontonn.Module.to("meta")directlyLayerInferenceEngine—meta_evictionhas one andlayer_inferencedoes not, so a quantized layer evicted throughlayer_inferencemay hit an overridden.to()layer_inference_test.py::TestModuleHelpersandmeta_eviction_test.pyboth green against the single implementationRelated: #1946, #1952, #13162, PR #13402