support arbitrary metric logging from torchmetrics

[sichu2023]

Description

Training and validation torchmetrics.Metric can is organized by TorchmetricsConfig. This encapsulates metric class instantiation and naming through get_metric_name instead of using field_factory.

Currently model parallelism is not supported and will raise NotImplementedError.

Type of changes

• New feature (non-breaking change which adds functionality)

[sichu2023]

/build-ci

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 90.16393% with 6 lines in your changes missing coverage. Please review.

Please upload report for BASE (main@9cec09f). Learn more about missing BASE report.

✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
...-packages/bionemo-llm/src/bionemo/llm/lightning.py	84.00%	4 Missing ⚠️
...emo-esm2/src/bionemo/esm2/scripts/finetune_esm2.py	85.71%	1 Missing ⚠️
...ionemo-esm2/src/bionemo/esm2/scripts/train_esm2.py	85.71%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #677   +/-   ##
=======================================
  Coverage        ?   86.27%           
=======================================
  Files           ?      119           
  Lines           ?     7249           
  Branches        ?        0           
=======================================
  Hits            ?     6254           
  Misses          ?      995           
  Partials        ?        0           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[farhadrgh]

I tentatively approve to unblock, but the tests are failing, and I wasn’t able to experiment with it.

[dorotat-nv]

ok, so the logic is that at least one rank will throw this error if model parallelism is initiated? Is it any metric logging? or only torch metric?

[sichu2023]

Yes and this is an imperfect solution which might lock the other devices.

We can condition on tensor/pipeline model parallelism If parallel_state has the functionality to inspect the tensor/pipeline model parallel size even when tp=pp=1. get_tensor_model_parallel_world_size and get_pipeline_model_parallel_world_size will throw an error in such case.

[sichu2023]

Would be great to have suggestions on alternatives.

[dorotat-nv]

imho, this error should be thrown when the metric config is initialised. In pydantic it could be

from pydantic import BaseModel, Field
 
class TorchmetricsConfig(BaseModel):
    class_path: str
    task: Literal["lm", "classification", "regression"]
    metric_name: str
    kwargs: Optional[dict[str, Any]] = Field(default_factory=dict)
    model_parallelism: bool = False
....
    def validate_parallel_logging(self) -> None:
        if self.model_parallelism:
            raise NotImplementedError(f"{self.__class__.__name__} logging is not implemented with model parallelism yet.")

and then you would need to intilised metrcis with

valid_metric = TorchmetricsConfig(
            class_path="Accuracy",
            task="classification",
            kwargs={
                "task": "multiclass",
                "threshold": 0.5,
                "num_classes": data_module.train_dataset.label_tokenizer.vocab_size,
            },
            metric_name="val_acc",
            model_parallelism: any([p > 1 [tensor_model_parallel, pipeline_model_parallel]])
        )

alternatively you can add parallelism check to the config and pass pp and tp settings

[sichu2023]

Since model parallelism is supported at BionemoLightningModule level I prefer to move the model parallelism check outside of TorchmetricsConfig. I have moved that into the training and finetuning scripts.

[pstjohn]

Could we get some tests for these new classes? Would be great to have unit tests around MetricConfig that show how it's used, as well as a short training run with some simple model that ensures it gets serialized correctly, produces the right results, etc.

[sichu2023]

@farhadrgh will have to pick your brain.

[farhadrgh]

The tests for fine-tuning are on a dummy dataset which overfits. I can replace that with a more reasonable dataset and uncomment these lines later

[sichu2023]

Thanks! Feel free to push your commits directly.

[sichu2023]

@farhadrgh will have to pick your brain.

[sichu2023]

@farhadrgh will have to pick your brain.

[dorotat-nv]

why not to use pydantic?

[dorotat-nv]

I would reimplement it to pydantic and add some validations

ie

from pydantic import BaseModel, Field, model_validator

class TorchmetricsConfig(BaseModel):
    class_path: str
    task: Literal["lm", "classification", "regression"]
    metric_name: str
    kwargs: Optional[dict[str, Any]] = Field(default_factory=dict)
    
    @model_validator(mode='after')
    def validate_class_path(self) -> 'TorchmetricsConfig':
        try:
            self.get_instance()
        except (ImportError, AttributeError) as e:
            raise ValueError(f"Invalid class_path: {self.class_path}. Error: {str(e)}")
        return self
    
    def get_instance(self) -> torchmetrics.Metric:
        """Dynamically imports and instantiates the metric class."""
        if "." in self.class_path:
            module_path, class_name = self.class_path.rsplit(".", 1)
            module = importlib.import_module(f"torchmetrics.{module_path}")
        else:
            class_name = self.class_path
            module = importlib.import_module("torchmetrics")
            
        cls = getattr(module, class_name)
        return cls(**self.kwargs)

[sichu2023]

Because we still have implementation beyond pydantic, e.g. train_esm.py and finetune_esm.py where most unittest and dev happen. That alone will be another discussion.

[dorotat-nv]

imho, this error should be thrown when the metric config is initialised. In pydantic it could be

from pydantic import BaseModel, Field
 
class TorchmetricsConfig(BaseModel):
    class_path: str
    task: Literal["lm", "classification", "regression"]
    metric_name: str
    kwargs: Optional[dict[str, Any]] = Field(default_factory=dict)
    model_parallelism: bool = False
....
    def validate_parallel_logging(self) -> None:
        if self.model_parallelism:
            raise NotImplementedError(f"{self.__class__.__name__} logging is not implemented with model parallelism yet.")

and then you would need to intilised metrcis with

valid_metric = TorchmetricsConfig(
            class_path="Accuracy",
            task="classification",
            kwargs={
                "task": "multiclass",
                "threshold": 0.5,
                "num_classes": data_module.train_dataset.label_tokenizer.vocab_size,
            },
            metric_name="val_acc",
            model_parallelism: any([p > 1 [tensor_model_parallel, pipeline_model_parallel]])
        )

alternatively you can add parallelism check to the config and pass pp and tp settings