INNER CODE UNIT · Python

p_trainable_params

Guitaricet/relora · torchrun_main.py:605

    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],
            output_device=local_rank,
        )

View source record →