wehub-resource-sync
513 lines
22 KiB
Python
513 lines
22 KiB
Python
# coding: utf-8
|
|
|
|
import functools
|
|
import os
|
|
import os.path as osp
|
|
import warnings
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.distributed.fsdp._traversal_utils as traversal_utils
|
|
from torch.distributed.device_mesh import init_device_mesh
|
|
from torch.distributed.fsdp import (
|
|
CPUOffload,
|
|
FullyShardedDataParallel as FSDP,
|
|
MixedPrecision,
|
|
BackwardPrefetch,
|
|
ShardingStrategy,
|
|
FullStateDictConfig,
|
|
StateDictType,
|
|
ShardedStateDictConfig,
|
|
ShardedOptimStateDictConfig
|
|
)
|
|
from torch.distributed.checkpoint import save as dcp_save, load as dcp_load
|
|
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
|
|
from safetensors.torch import load_file
|
|
|
|
from modeling.lance.modeling_utils import MLPconnector, TimestepEmbedder
|
|
from modeling.lance.modeling_utils import PositionEmbedding3D
|
|
|
|
from modeling.lance.qwen2_navit import (
|
|
# Qwen2ForCausalLM,
|
|
Qwen2DecoderLayer,
|
|
Qwen2MoEDecoderLayer,
|
|
Qwen2MoTDecoderLayer,
|
|
)
|
|
from common.utils.fs import mkdir, is_hdfs_path, copy, exists
|
|
from common.utils.save import get_local_dir, save
|
|
|
|
# Ignore the specific FutureWarning at module import time
|
|
warnings.filterwarnings(
|
|
"ignore",
|
|
category=FutureWarning,
|
|
module="torch.distributed.fsdp.fully_sharded_data_parallel"
|
|
)
|
|
|
|
|
|
# -------------------------- helpers --------------------------
|
|
def _rank() -> int:
|
|
return dist.get_rank() if dist.is_available() and dist.is_initialized() else 0
|
|
|
|
|
|
def _world() -> int:
|
|
return dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1
|
|
|
|
|
|
def __barrier__():
|
|
if dist.is_available() and dist.is_initialized():
|
|
dist.barrier()
|
|
|
|
|
|
def _copy_model_metadata_files(source_model_path, save_path, blocking=True):
|
|
if not source_model_path or _rank() != 0:
|
|
return
|
|
|
|
metadata_filenames = [
|
|
"generation_config.json",
|
|
"llm_config.json",
|
|
"merges.txt",
|
|
"tokenizer.json",
|
|
"tokenizer_config.json",
|
|
"vocab.json",
|
|
]
|
|
for filename in metadata_filenames:
|
|
src = osp.join(source_model_path, filename)
|
|
if exists(src):
|
|
copy(src, osp.join(save_path, filename), blocking=blocking)
|
|
|
|
|
|
def _save_load_report(report_dir, report_name, title, msg, logger):
|
|
missing = list(getattr(msg, "missing_keys", []))
|
|
unexpected = list(getattr(msg, "unexpected_keys", []))
|
|
logger.info(f"{title}: missing={len(missing)}, unexpected={len(unexpected)}")
|
|
if report_dir is None or _rank() != 0:
|
|
return
|
|
|
|
mkdir(report_dir)
|
|
report_path = osp.join(report_dir, report_name)
|
|
lines = [
|
|
title,
|
|
f"missing_count: {len(missing)}",
|
|
f"unexpected_count: {len(unexpected)}",
|
|
"",
|
|
"missing_keys:",
|
|
*[str(item) for item in missing],
|
|
"",
|
|
"unexpected_keys:",
|
|
*[str(item) for item in unexpected],
|
|
]
|
|
with open(report_path, "w", encoding="utf-8") as f:
|
|
f.write("\n".join(lines))
|
|
f.write("\n")
|
|
logger.info(f"Saved checkpoint load report to {report_path}")
|
|
|
|
|
|
class FSDPConfig:
|
|
def __init__(
|
|
self,
|
|
sharding_strategy,
|
|
backward_prefetch,
|
|
cpu_offload,
|
|
num_replicate,
|
|
num_shard=8,
|
|
use_orig_params=False,
|
|
):
|
|
self.sharding_strategy = sharding_strategy
|
|
self.backward_prefetch = backward_prefetch
|
|
self.cpu_offload = cpu_offload
|
|
self.num_replicate = num_replicate
|
|
self.num_shard = num_shard
|
|
self.use_orig_params = use_orig_params
|
|
|
|
|
|
def fsdp_wrapper(original_model, fsdp_config: FSDPConfig, ignored_modules=[], mixed_precision_override=None):
|
|
if fsdp_config.sharding_strategy == 'HYBRID_SHARD':
|
|
device_mesh = init_device_mesh(
|
|
"cuda",
|
|
mesh_shape=(fsdp_config.num_replicate, fsdp_config.num_shard),
|
|
mesh_dim_names=("replicate", "shard")
|
|
)
|
|
else:
|
|
device_mesh = None
|
|
|
|
mp = mixed_precision_override or MixedPrecision(
|
|
param_dtype=torch.bfloat16,
|
|
reduce_dtype=torch.bfloat16,
|
|
# reduce_dtype=torch.float32, # TODO: using torch.float32 converts bfloat16 to float32
|
|
buffer_dtype=torch.bfloat16,
|
|
)
|
|
|
|
return FSDP(
|
|
original_model,
|
|
auto_wrap_policy=functools.partial(
|
|
transformer_auto_wrap_policy,
|
|
transformer_layer_cls={
|
|
# Qwen2ForCausalLM,
|
|
Qwen2DecoderLayer,
|
|
Qwen2MoEDecoderLayer,
|
|
Qwen2MoTDecoderLayer,
|
|
MLPconnector,
|
|
TimestepEmbedder,
|
|
# PositionEmbedding, # NOTE: NOT USED
|
|
PositionEmbedding3D,
|
|
},
|
|
),
|
|
ignored_modules=ignored_modules,
|
|
mixed_precision=mp,
|
|
device_id=dist.get_rank() % torch.cuda.device_count(),
|
|
sharding_strategy=ShardingStrategy[fsdp_config.sharding_strategy],
|
|
backward_prefetch=BackwardPrefetch[fsdp_config.backward_prefetch],
|
|
cpu_offload=CPUOffload(offload_params=fsdp_config.cpu_offload),
|
|
device_mesh=device_mesh,
|
|
use_orig_params=fsdp_config.use_orig_params,
|
|
)
|
|
|
|
|
|
class FSDPCheckpoint:
|
|
@staticmethod
|
|
def fsdp_save_fsdp_ckpt(ckpt_dir, train_steps, model, ema_model, optimizer, scheduler, data_status, logger, fsdp_config, blocking=True, **kwargs):
|
|
save_path = osp.join(ckpt_dir, f"{train_steps:07d}")
|
|
mkdir(save_path)
|
|
logger.info(f"Begin saving checkpoint info to {save_path}")
|
|
source_model_path = kwargs.get("source_model_path")
|
|
|
|
local_save_dir = get_local_dir() if is_hdfs_path(save_path) else None
|
|
rank = _rank()
|
|
world = _world()
|
|
|
|
if fsdp_config.sharding_strategy == "HYBRID_SHARD":
|
|
assert world == fsdp_config.num_shard * fsdp_config.num_replicate, f"world={world} != shard({fsdp_config.num_shard})*replicate({fsdp_config.num_replicate})"
|
|
|
|
# ---- 0) For HDFS targets, write DCP output to a local temp directory first, then copy it back ----
|
|
dcp_root = save_path
|
|
if is_hdfs_path(save_path):
|
|
dcp_root = osp.join(get_local_dir(), osp.basename(save_path))
|
|
os.makedirs(dcp_root, exist_ok=True)
|
|
|
|
# ---- 1) save sharded via DCP ----
|
|
if kwargs.get("flag_save_shard_model", False): # Whether to save the sharded model
|
|
# ---- 1.1) EMA (sharded via DCP) ----
|
|
if ema_model is not None: # NOTE: ema_model is currently None
|
|
with FSDP.state_dict_type(
|
|
ema_model,
|
|
StateDictType.SHARDED_STATE_DICT,
|
|
ShardedStateDictConfig(offload_to_cpu=True),
|
|
):
|
|
ema_model_state = ema_model.state_dict()
|
|
dcp_save(ema_model_state, checkpoint_id=osp.join(dcp_root, "ema"))
|
|
del ema_model_state
|
|
import gc; gc.collect(); torch.cuda.empty_cache()
|
|
|
|
__barrier__()
|
|
|
|
# ---- 1.2) Model (sharded via DCP) ----
|
|
with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT, ShardedStateDictConfig(offload_to_cpu=True)):
|
|
model_state = model.state_dict()
|
|
dcp_save(model_state, checkpoint_id=os.path.join(dcp_root, "model"))
|
|
del model_state
|
|
import gc; gc.collect(); torch.cuda.empty_cache()
|
|
|
|
__barrier__()
|
|
|
|
# ---- 2) Model FULL ----
|
|
if kwargs.get("flag_save_full_model", True): # Whether to save the full model
|
|
# ---- 2.1) EMA Model FULL ----
|
|
if ema_model is not None:
|
|
with FSDP.state_dict_type(ema_model, StateDictType.FULL_STATE_DICT, FullStateDictConfig(rank0_only=True, offload_to_cpu=True)):
|
|
sd = ema_model.state_dict()
|
|
|
|
__barrier__() # Synchronize once inside the FULL context
|
|
|
|
if rank == 0:
|
|
# Optional contiguous conversion
|
|
for k, v in list(sd.items()):
|
|
if isinstance(v, torch.Tensor) and not v.is_contiguous():
|
|
sd[k] = v.contiguous()
|
|
save(sd, osp.join(save_path, "ema.safetensors"), blocking=blocking, local_dir=local_save_dir)
|
|
|
|
__barrier__()
|
|
|
|
del sd
|
|
import gc; gc.collect(); torch.cuda.empty_cache()
|
|
|
|
# ---- 2.2) Model FULL ----
|
|
# NOTE: Saved model fine-tuning loss was verified on 2025-08-21
|
|
with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, FullStateDictConfig(rank0_only=True, offload_to_cpu=True)):
|
|
sd = model.state_dict() # Even with rank0-only, all ranks must call state_dict() for collective communication
|
|
|
|
__barrier__() # Synchronize once inside the FULL context
|
|
|
|
if rank == 0:
|
|
# Optional contiguous conversion
|
|
for k, v in list(sd.items()):
|
|
if isinstance(v, torch.Tensor) and not v.is_contiguous():
|
|
sd[k] = v.contiguous()
|
|
save(sd, osp.join(save_path, "model.safetensors"), blocking=blocking, local_dir=local_save_dir)
|
|
|
|
__barrier__() # Synchronize inside the FULL context so non-rank0 processes do not exit too early
|
|
|
|
del sd
|
|
import gc; gc.collect(); torch.cuda.empty_cache()
|
|
|
|
_copy_model_metadata_files(source_model_path, save_path, blocking=blocking)
|
|
__barrier__()
|
|
|
|
# --- 3) If target is HDFS, copy sharded dirs up ---
|
|
if is_hdfs_path(save_path): # NOTE: DCP output needs an extra copy step; skip if sharded checkpoints were not saved
|
|
for sub in ("model", "ema"):
|
|
src = osp.join(dcp_root, sub)
|
|
if os.path.exists(src):
|
|
copy(src, save_path, blocking=blocking) # fix: directory copy
|
|
logger.info(f"Copy {src} to HDFS or Local path: {save_path} done.")
|
|
|
|
# ---- 3) Optimizer (sharded via DCP) ----
|
|
if kwargs.get("flag_save_optimizer", True): # Whether to save the optimizer; enabled by default for full resume
|
|
# Determine shard file name (keeps your original convention)
|
|
if fsdp_config.sharding_strategy == "FULL_SHARD":
|
|
shard_index = rank
|
|
total_shards = world
|
|
elif fsdp_config.sharding_strategy == "HYBRID_SHARD":
|
|
shard_index = rank % fsdp_config.num_shard
|
|
total_shards = fsdp_config.num_shard
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
opt_path = osp.join(save_path, f"optimizer.{shard_index:05d}-of-{total_shards:05d}.pt")
|
|
|
|
# Export the *sharded* optimizer state dict (no rank0 aggregation; low memory)
|
|
with FSDP.state_dict_type(
|
|
model,
|
|
StateDictType.SHARDED_STATE_DICT,
|
|
optim_state_dict_config=ShardedOptimStateDictConfig(offload_to_cpu=True), # torch 2.5.1: no `group` arg
|
|
):
|
|
osd = FSDP.optim_state_dict(model, optimizer)
|
|
|
|
# Save the shard (HYBRID: only first `num_shard` ranks write, matching your pattern)
|
|
try:
|
|
if fsdp_config.sharding_strategy == "FULL_SHARD":
|
|
save(osd, opt_path, blocking=blocking, local_dir=local_save_dir)
|
|
elif fsdp_config.sharding_strategy == "HYBRID_SHARD":
|
|
if rank < fsdp_config.num_shard:
|
|
save(osd, opt_path, blocking=blocking, local_dir=local_save_dir)
|
|
finally:
|
|
del osd
|
|
import gc; gc.collect(); torch.cuda.empty_cache()
|
|
|
|
# ---- Scheduler (rank0) ----
|
|
if rank == 0 and scheduler is not None:
|
|
# torch.save(scheduler.state_dict(), osp.join(save_path, "scheduler.pt"))
|
|
save(scheduler.state_dict(), osp.join(save_path, "scheduler.pt"), blocking=blocking, local_dir=local_save_dir)
|
|
|
|
# ---- Data status (per-rank) ----
|
|
if rank == 0 and data_status is not None:
|
|
save(data_status, osp.join(save_path, "data_status.pt"), blocking=blocking, local_dir=local_save_dir)
|
|
|
|
del data_status
|
|
import gc; gc.collect(); torch.cuda.empty_cache()
|
|
|
|
__barrier__()
|
|
|
|
return
|
|
|
|
@staticmethod
|
|
def try_load_ckpt(resume_from, logger, model, ema_model=None, resume_from_ema=False, report_dir=None):
|
|
# TODO: verify this
|
|
if resume_from is not None and osp.exists(resume_from):
|
|
logger.info(f"Loading checkpoint from {resume_from}.")
|
|
if resume_from_ema:
|
|
model_state_dict_path = osp.join(resume_from, f"ema.safetensors")
|
|
else:
|
|
model_state_dict_path = osp.join(resume_from, f"model.safetensors")
|
|
model_state_dict = load_file(model_state_dict_path, device="cpu")
|
|
# NOTE position embeds are fixed sinusoidal embeddings, so we can just pop it off,
|
|
# which makes it easier to adapt to different resolutions.
|
|
for key in ["latent_pos_embed.pos_embed", "vit_pos_embed.pos_embed"]:
|
|
if key in model_state_dict:
|
|
model_state_dict.pop(key)
|
|
msg = model.load_state_dict(model_state_dict, strict=False)
|
|
_save_load_report(report_dir, "resume_checkpoint_load_report_model.txt", f"Resume model checkpoint: {model_state_dict_path}", msg, logger)
|
|
del model_state_dict
|
|
|
|
if ema_model is not None:
|
|
ema_state_dict_path = osp.join(resume_from, f"ema.safetensors")
|
|
if not osp.exists(ema_state_dict_path):
|
|
logger.info(f"Replicating EMA model from {model_state_dict_path}.")
|
|
ema_state_dict_path = model_state_dict_path
|
|
ema_state_dict = load_file(ema_state_dict_path, device="cpu")
|
|
# NOTE position embeds are fixed sinusoidal embeddings, so we can just pop it off,
|
|
# which makes it easier to adapt to different resolutions.
|
|
for key in ["latent_pos_embed.pos_embed", "vit_pos_embed.pos_embed"]:
|
|
if key in ema_state_dict:
|
|
ema_state_dict.pop(key)
|
|
msg = ema_model.load_state_dict(ema_state_dict, strict=False)
|
|
_save_load_report(report_dir, "resume_checkpoint_load_report_ema.txt", f"Resume EMA checkpoint: {ema_state_dict_path}", msg, logger)
|
|
del ema_state_dict
|
|
else:
|
|
logger.info(f"Training from scratch.")
|
|
return model, ema_model
|
|
|
|
@staticmethod
|
|
def try_load_fsdp_ckpt(resume_from, logger, model, ema_model=None, resume_from_ema=False, report_dir=None):
|
|
# TODO: verify this
|
|
if resume_from is None or not os.path.exists(resume_from):
|
|
logger.info("Training from scratch.")
|
|
return model, ema_model
|
|
|
|
logger.info(f"Loading checkpoint from {resume_from}.")
|
|
|
|
# ---- Model (or EMA) via DCP ----
|
|
load_dir = osp.join(resume_from, "ema" if resume_from_ema else "model")
|
|
assert isinstance(model, FSDP)
|
|
with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT, ShardedStateDictConfig(offload_to_cpu=True)):
|
|
model_state = model.state_dict()
|
|
dcp_load(model_state, checkpoint_id=load_dir)
|
|
for k in ["latent_pos_embed.pos_embed", "vit_pos_embed.pos_embed"]:
|
|
if k in model_state:
|
|
model_state.pop(k)
|
|
msg = model.load_state_dict(model_state, strict=False)
|
|
_save_load_report(report_dir, "resume_fsdp_checkpoint_load_report_model.txt", f"Resume FSDP model checkpoint: {load_dir}", msg, logger)
|
|
del model_state
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
# ---- EMA (optional) ----
|
|
if ema_model is not None:
|
|
ema_dir = osp.join(resume_from, "ema")
|
|
assert isinstance(ema_model, FSDP)
|
|
with FSDP.state_dict_type(ema_model, StateDictType.SHARDED_STATE_DICT, ShardedStateDictConfig(offload_to_cpu=True)):
|
|
ema_state = ema_model.state_dict()
|
|
dcp_load(ema_state, checkpoint_id=ema_dir)
|
|
for k in ["latent_pos_embed.pos_embed", "vit_pos_embed.pos_embed"]:
|
|
if k in ema_state:
|
|
ema_state.pop(k)
|
|
msg = ema_model.load_state_dict(ema_state, strict=False)
|
|
_save_load_report(report_dir, "resume_fsdp_checkpoint_load_report_ema.txt", f"Resume FSDP EMA checkpoint: {ema_dir}", msg, logger)
|
|
del ema_state
|
|
import gc; gc.collect(); torch.cuda.empty_cache()
|
|
|
|
return model, ema_model
|
|
|
|
@staticmethod
|
|
def try_load_train_state(resume_from, optimizer, scheduler, fsdp_config):
|
|
if resume_from is not None and osp.exists(resume_from):
|
|
if fsdp_config.sharding_strategy == "FULL_SHARD":
|
|
shard_index = dist.get_rank()
|
|
total_shards = dist.get_world_size()
|
|
elif fsdp_config.sharding_strategy == "HYBRID_SHARD":
|
|
shard_index = dist.get_rank() % fsdp_config.num_shard
|
|
total_shards = fsdp_config.num_shard
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
optimizer_state_dict_path = osp.join(resume_from, f"optimizer.{shard_index:05d}-of-{total_shards:05d}.pt")
|
|
optimizer_state_dict = torch.load(optimizer_state_dict_path, map_location="cpu", weights_only=True)
|
|
optimizer.load_state_dict(optimizer_state_dict)
|
|
del optimizer_state_dict
|
|
|
|
scheduler_state_dict_path = osp.join(resume_from, "scheduler.pt")
|
|
scheduler_state_dict = torch.load(scheduler_state_dict_path, weights_only=True, map_location="cpu")
|
|
scheduler.load_state_dict(scheduler_state_dict)
|
|
del scheduler_state_dict
|
|
|
|
train_steps = int(osp.basename(osp.normpath(resume_from))) + 1
|
|
"""
|
|
data_status = [
|
|
{
|
|
dataset_name: {
|
|
worker_id: [parquet_idx, row_group_id, row_idx],
|
|
},
|
|
},
|
|
]
|
|
"""
|
|
data_status_path = osp.join(resume_from, "data_status.pt")
|
|
if osp.exists(data_status_path):
|
|
data_status = torch.load(data_status_path, weights_only=True, map_location="cpu")
|
|
local_rank = dist.get_rank()
|
|
if local_rank < len(data_status):
|
|
data_status = data_status[local_rank]
|
|
else:
|
|
data_status = None
|
|
else:
|
|
data_status = None
|
|
else:
|
|
train_steps = 0
|
|
data_status = None
|
|
return optimizer, scheduler, train_steps, data_status
|
|
|
|
|
|
def grad_checkpoint_check_fn(module):
|
|
module_options = (
|
|
Qwen2DecoderLayer,
|
|
Qwen2MoEDecoderLayer,
|
|
Qwen2MoTDecoderLayer,
|
|
MLPconnector,
|
|
)
|
|
return isinstance(module, module_options)
|
|
|
|
|
|
def fsdp_ema_setup(ema_model, fsdp_config, ignored_modules=[]):
|
|
ema_model.eval()
|
|
for param in ema_model.parameters():
|
|
param.requires_grad = False
|
|
|
|
ema_model = fsdp_wrapper(ema_model, fsdp_config, ignored_modules=ignored_modules)
|
|
return ema_model
|
|
|
|
|
|
@torch.no_grad()
|
|
def fsdp_ema_update(ema_model, model, decay=0.9999):
|
|
ema_handles = traversal_utils._get_fsdp_handles(ema_model)
|
|
new_handles = traversal_utils._get_fsdp_handles(model)
|
|
assert len(ema_handles) == len(new_handles)
|
|
ema_params = []
|
|
new_params = []
|
|
|
|
for ema_handle, new_handle in zip(ema_handles, new_handles):
|
|
if ema_handle.flat_param is not None and new_handle.flat_param.requires_grad:
|
|
ema_params.append(ema_handle.flat_param.data)
|
|
new_params.append(new_handle.flat_param.data.to(dtype=ema_handle.flat_param.dtype))
|
|
|
|
torch._foreach_mul_(ema_params, decay)
|
|
torch._foreach_add_(ema_params, new_params, alpha=1 - decay)
|
|
|
|
|
|
# =============================================== CPU EMA implementation, pending verification =====================================================
|
|
def fsdp_ema_setup_v2(ema_model, fsdp_config, ignored_modules=[], backend="fsdp_cpu_offload"):
|
|
ema_model.eval()
|
|
for param in ema_model.parameters():
|
|
param.requires_grad = False
|
|
|
|
if backend == "none_cpu_plain":
|
|
return ema_model.cpu()
|
|
|
|
# Default: wrap with FSDP, but force CPU offload so GPU memory use is near zero
|
|
ema_cfg = FSDPConfig(
|
|
sharding_strategy=fsdp_config.sharding_strategy,
|
|
backward_prefetch=fsdp_config.backward_prefetch,
|
|
cpu_offload=True, # Key setting
|
|
num_replicate=fsdp_config.num_replicate,
|
|
num_shard=fsdp_config.num_shard,
|
|
use_orig_params=fsdp_config.use_orig_params,
|
|
)
|
|
|
|
mp_ema = MixedPrecision(param_dtype=torch.float32, reduce_dtype=torch.bfloat16, buffer_dtype=torch.bfloat16)
|
|
ema_model = fsdp_wrapper(ema_model, ema_cfg, ignored_modules=ignored_modules, mixed_precision_override=mp_ema)
|
|
return ema_model
|
|
|
|
|
|
@torch.no_grad()
|
|
def fsdp_ema_update_v2(ema_model, model, decay=0.9999):
|
|
ema_handles = traversal_utils._get_fsdp_handles(ema_model)
|
|
new_handles = traversal_utils._get_fsdp_handles(model)
|
|
assert len(ema_handles) == len(new_handles)
|
|
ema_params = []
|
|
new_params = []
|
|
|
|
for ema_handle, new_handle in zip(ema_handles, new_handles):
|
|
if ema_handle.flat_param is not None and new_handle.flat_param.requires_grad:
|
|
# EMA stays on CPU: ensure new_params are also on CPU before foreach ops
|
|
ema_params.append(ema_handle.flat_param.data) # CPU
|
|
new_params.append(new_handle.flat_param.data.to(device="cpu", dtype=ema_handle.flat_param.dtype, non_blocking=True))
|
|
|
|
torch._foreach_mul_(ema_params, decay)
|
|
torch._foreach_add_(ema_params, new_params, alpha=1 - decay)
|