mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 02:06:15 +00:00
fix(ema): default EMAConfig.enable to False (opt-in)
EMA was on by default, so every training run on the branch (incl. VLA-JEPA and other non-flow-matching policies) created a full fp32 shadow copy. EMA only benefits flow-matching/diffusion policies (pi0/pi05/pi052). Make it opt-in via --ema.enable=true; the pi05/pi052 recipes already pass that flag. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -98,14 +98,14 @@ class EMAConfig:
|
|||||||
Cost: 1× model params in fp32 shadow (~13 GB for pi052's 3.3B
|
Cost: 1× model params in fp32 shadow (~13 GB for pi052's 3.3B
|
||||||
params) + one elementwise update per training step (~1% step time).
|
params) + one elementwise update per training step (~1% step time).
|
||||||
|
|
||||||
On by default — matches openpi (JAX) which ships EMA on for every
|
Off by default (opt-in): EMA is only beneficial for flow-matching /
|
||||||
config, and closes the gap with the openpi PyTorch port which
|
diffusion policies (pi0/pi05/pi052), and the fp32 shadow copy is pure
|
||||||
explicitly lists EMA as unsupported. Set ``--ema.enable=false`` to
|
overhead for other policies (e.g. VLA-JEPA). Set ``--ema.enable=true``
|
||||||
disable for short runs / memory-constrained training where the
|
to turn it on (the pi05/pi052 training recipes do this). openpi (JAX)
|
||||||
extra fp32 shadow copy is the bottleneck.
|
ships EMA on for every config; enable it explicitly to match that.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
enable: bool = True
|
enable: bool = False
|
||||||
# Target EMA decay β in θ_ema ← β·θ_ema + (1-β)·θ_live (passed to
|
# Target EMA decay β in θ_ema ← β·θ_ema + (1-β)·θ_live (passed to
|
||||||
# ema-pytorch as ``beta``).
|
# ema-pytorch as ``beta``).
|
||||||
# 0.999 — last ~1000 steps; pi05_libero default in openpi
|
# 0.999 — last ~1000 steps; pi05_libero default in openpi
|
||||||
|
|||||||
Reference in New Issue
Block a user