INNER CODE UNIT · Python
n_total_params
Guitaricet/relora · torchrun_main.py:603
n_total_params = sum(p.numel() for p in model.parameters())
n_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
p_trainable_params = n_trainable_params / n_total_params
# ##############################
# Distributed wrapping
if args.distributed_type == "fsdp":
logger.info("Wrapping model with FSDP")
raise RuntimeError("FSDP is not supported anymore. "
"There were a lot of isses with ReLoRA and FSDP "
"and no speed or memory improvements.")
model: Union[FSDP, ReLoRaModel, LlamaForCausalLM] = training_utils.initialize_fsdp(model, dtype=args.dtype)
elif args.distributed_type == "ddp":
logger.info("Wrapping model with DDP")
model: Union[ReLoRaModel, LlamaForCausalLM] = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[local_rank],