Method __init__

axlearn/common/trainer.py:246–381

(
        self,
        cfg: Config,
        *,
        parent: Optional[Module],
        devices: Optional[np.ndarray] = None,
    )

        # and compilation process every train step and rely on compilation cache to prevent
        # excessive recompilations. Note: this could introduce overhead to training due to
        # pre-compilation checks (such as sharding check) that increases the step time for some
        # models. Note that this cache is always disabled at steps when xsc is enabled.
        # Defaults to None which is interpreted as True.
        cache_compiled_train_step: Optional[bool] = None

        # Log the loss value every n steps. Defaults to None which is interpreted as every
        # 100 steps.
        log_every_n_steps: Optional[int] = None

    def __init__(
        self,
        cfg: Config,
        *,
        parent: Optional[Module],
        devices: Optional[np.ndarray] = None,
    ):
        super().__init__(cfg, parent=parent)
        cfg = self.config

        if not cfg.prune_empty_state_updates:
            raise ValueError(
                "Setting prune_empty_state_updates to False is no longer supported.\n"
                "The config option will be removed in a future AXLearn version."
            )

        self._step: int = None
        self._trainer_state: TrainerState = None
        self._jit_train_step: jax.stages.Wrapped = None
        self._watchdog_stopping = None
        self._watchdog_thread = None
        self._device_monitor = maybe_instantiate(cfg.device_monitor)
        self._recorder = maybe_instantiate(cfg.recorder)
        self._is_initialized: bool = False
        self._maybe_record_event(measurement.Event.START_ACCELERATOR_INIT)

        if cfg.model.dtype is None:
            raise ValueError(f"dtype must be explicitly specified for {self.path()}.model")
        if cfg.model.param_init is None:
            cfg.model.param_init = DefaultInitializer.default_config()
            logging.info(
                "model.param_init is not specified. Default to DefaultInitializer: %s",
                cfg.model.param_init,
            )

        self._per_param_train_dtype = maybe_instantiate(
            canonicalize_per_param_dtype(cfg.train_dtype)
        )

        # Create the device mesh.
        if devices is None:
            self._step_log(
                "Devices: global=%s local=%s %s",
                jax.device_count(),
                jax.local_device_count(),
                [device.platform for device in jax.local_devices()],
            )
        else:
            local_devices = [d for d in devices.flatten() if d.process_index == jax.process_index()]