Repository: oumi-ai/oumi
File: src/oumi/utils/torch_utils.py (Python)

Code before:
@dataclass(frozen=True)
class ModelParameterCount:
    all_params: int
    trainable_params: int
    embedding_params: int

    def __post_init__(self):
        """Ensure that the parameters are valid."""
        for name, value in [
            ("all_params", self.all_params),
            ("trainable_params", self.trainable_params),
            ("embedding_params", self.embedding_params),
        ]:
            if value < 0:
                raise ValueError(f"`{name}` must be >= 0.")
        if self.trainable_params > self.all_params:
            raise ValueError("`trainable_params` cannot be greater than `all_params`.")
        if self.embedding_params > self.all_params:
            raise ValueError("`embedding_params` cannot be greater than `all_params`.")

    @property
    def trainable_params_percent(self) -> float:
        """Percentage of trainable parameters [0.0, 100.0], rounded to 2 decimals."""
        if self.all_params == 0:
            return 0.0
        return round(100 * self.trainable_params / self.all_params, 2)

    @property
    def frozen_params_percent(self) -> float:
        """Percentage of frozen parameters [0.0, 100.0], rounded to 2 decimals."""
        return round(100.0 - self.trainable_params_percent, 2)


def _get_parameter_names(
    model: torch.nn.Module, forbidden_layer_types: list[Any]
) -> list[str]:
    """Returns the names of the model parameters that are not inside a forbidden layer.

    Borrowed from
    https://github.com/huggingface/transformers/blob/main/src/transformers/trainer.py.
    """
    result = []
    for name, child in model.named_children():
        result += [
            f"{name}.{n}"
            for n in _get_parameter_names(child, forbidden_layer_types)
            if not isinstance(child, tuple(forbidden_layer_types))
        ]
    # Add model specific parameters (defined with nn.Parameter) since they are not in
    # any child.
    result += list(model._parameters.keys())

Review comment (suggestion):
I don't think we should  be `round()`-ing here . I suggest to return unrounded ratio.

We can round as needed e.g., in formatting function:  `{params.frozen_params_percent:.2f}%`

Code after:
@dataclass(frozen=True)
class ModelParameterCount:
    all_params: int
    trainable_params: int
    embedding_params: int

    def __post_init__(self):
        """Ensure that the parameters are valid."""
        for name, value in [
            ("all_params", self.all_params),
            ("trainable_params", self.trainable_params),
            ("embedding_params", self.embedding_params),
        ]:
            if value < 0:
                raise ValueError(f"`{name}` ({value}) must be >= 0.")
            if self.trainable_params > self.all_params:
                raise ValueError(
                    f"`trainable_params` ({self.trainable_params}) cannot be "
                    f"greater than `all_params` ({self.all_params})."
                )
            if self.embedding_params > self.all_params:
                raise ValueError(
                    f"`embedding_params` ({self.embedding_params}) cannot be "
                    f"greater than `all_params` ({self.all_params})."
                )

    @property
    def trainable_params_percent(self) -> float:
        """Percentage of trainable parameters [0.0, 100.0]."""
        if self.all_params == 0:
            return 0.0
        return 100 * self.trainable_params / self.all_params

    @property
    def frozen_params_percent(self) -> float:
        """Percentage of frozen parameters [0.0, 100.0]."""
        return 100.0 - self.trainable_params_percent


def _get_parameter_names(
    model: torch.nn.Module, forbidden_layer_types: list[Any]
) -> list[str]:
    """Returns the names of the model parameters that are not inside a forbidden layer.

    Borrowed from
    https://github.com/huggingface/transformers/blob/main/src/transformers/trainer.py.
    """
    result = []
    for name, child in model.named_children():
        result += [
            f"{name}.{n}"