Repository: Physical-Intelligence/openpi
File: src/openpi/training/config.py (Python)

Code before:
    # Global batch size.
    batch_size: int = 32
    # Number of workers to use for the data loader.
    num_workers: int = 2
    # Number of train steps (batches) to run.
    num_train_steps: int = 30_000

    # How often to log training metrics.
    log_interval: int = 100
    # How often to save checkpoints.
    save_interval: int = 1000
    # How often to keep checkpoints.
    keep_interval: int = 5000

    # If true, will overwrite the checkpoint directory if it already exists.
    overwrite: bool = False
    # If true, will resume training from the last checkpoint.
    resume: bool = False

    # Keyword arguments to pass to the policy's sample method.
    sample_kwargs: dict[str, Any] | None = None

    # If true, will enable wandb logging.
    wandb_enabled: bool = True

    # If the value is greater than 1, FSDP will be enabled and shard across number of specified devices.
    # eg. if total device is 4 and fsdp parallel devices is 2; then the model will shard to 2 devices and run
    # data parallel between 2 groups.
    num_fsdp_devices: int = 1

    @property
    def metadata_dir(self) -> pathlib.Path:
        """Get the metadata directory for this config."""
        return (pathlib.Path(self.metadata_base_dir) / self.name).resolve()

    @property
    def checkpoint_dir(self) -> pathlib.Path:
        """Get the checkpoint directory for this config."""
        if not self.exp_name:
            raise ValueError("--exp_name must be set")
        return (pathlib.Path(self.checkpoint_base_dir) / self.name / self.exp_name).resolve()

    def create_model(self) -> _model.Model:
        """Create a model for this config."""
        return _model.Model(
            module=self.module,
            action_dim=self.action_dim,
            action_horizon=self.action_horizon,
            max_token_len=self.max_token_len,
        )


Review comment (performance):
Just in case, can you clarify that this will reduce the overall memory consumption but may make training slower?

Code after:
    # Global batch size.
    batch_size: int = 32
    # Number of workers to use for the data loader.
    num_workers: int = 2
    # Number of train steps (batches) to run.
    num_train_steps: int = 30_000

    # How often to log training metrics.
    log_interval: int = 100
    # How often to save checkpoints.
    save_interval: int = 1000
    # How often to keep checkpoints.
    keep_interval: int = 5000

    # If true, will overwrite the checkpoint directory if it already exists.
    overwrite: bool = False
    # If true, will resume training from the last checkpoint.
    resume: bool = False

    # Keyword arguments to pass to the policy's sample method.
    sample_kwargs: dict[str, Any] | None = None

    # If true, will enable wandb logging.
    wandb_enabled: bool = True

    # If the value is greater than 1, FSDP will be enabled and shard across number of specified devices; overall
    # device memory will be reduced but training could potentially be slower.
    # eg. if total device is 4 and fsdp devices is 2; then the model will shard to 2 devices and run
    # data parallel between 2 groups of devices.
    fsdp_devices: int = 1

    @property
    def metadata_dir(self) -> pathlib.Path:
        """Get the metadata directory for this config."""
        return (pathlib.Path(self.metadata_base_dir) / self.name).resolve()

    @property
    def checkpoint_dir(self) -> pathlib.Path:
        """Get the checkpoint directory for this config."""
        if not self.exp_name:
            raise ValueError("--exp_name must be set")
        return (pathlib.Path(self.checkpoint_base_dir) / self.name / self.exp_name).resolve()

    def create_model(self) -> _model.Model:
        """Create a model for this config."""
        return _model.Model(
            module=self.module,
            action_dim=self.action_dim,
            action_horizon=self.action_horizon,
            max_token_len=self.max_token_len,
        )