llm.engine.initialize
initialize() ¶
initialize(
model: torch.nn.Module,
optimizer: torch.optim.Optimizer,
criterion: torch.nn.Module,
scheduler: torch.optim.lr_scheduler._LRScheduler,
accumulation_steps: int = 1,
dtype: torch.dtype | None = None,
max_norm: float | None = None,
**kwargs: Any
) -> tuple[
torch.nn.Module,
BaseOptimizer,
torch.nn.Module,
torch.optim.lr_scheduler._LRScheduler,
]
Enable advanced training features.
This method allows you to easily wrap your training objects with transparent wrappers that enable advanced training features.
Parameters:
-
model
(Module
) –Model being trained.
-
optimizer
(Optimizer
) –Training optimizer.
-
criterion
(Module
) –Training loss function.
-
scheduler
(_LRScheduler
) –LR scheduler.
-
accumulation_steps
(int
, default:1
) –Number of forward/backward passes between optimizer steps.
-
dtype
(dtype | None
, default:None
) –Optional data type for mixed precision training.
-
max_norm
(float | None
, default:None
) –Optional maximum norm of gradients to clip to.
-
kwargs
(Any
, default:{}
) –Keyword arguments to pass to the gradient scaler.
Returns:
-
tuple[Module, BaseOptimizer, Module, _LRScheduler]
–Tuple of the wrapped model, optimizer, loss, and scheduler.