Adds configuration field optimizer.record_update_metrics, which
defaults to False, but when set to True will trigger AdamW to
collect the step size norm and absolute max for each parameter.
Changes the behavior of the Lion optimizer to only record the update cosine
similarity when optimizer.record_update_metrics is True in order to be
consistent with the API.
optimizer.record_update_metrics
, which defaults toFalse
, but when set toTrue
will trigger AdamW to collect the step size norm and absolute max for each parameter.optimizer.record_update_metrics
isTrue
in order to be consistent with the API.See https://wandb.ai/ai2-llm/petew-update-logging?nw=nwuserepwalsh for a comparison of a run with and without optimizer update logging.