Basic implementation of WeightedCrossEntropy torchmetric. #1
Conversation
| @@ -835,13 +836,47 @@ def forward( | |||
| return hidden_states, all_hidden_states | |||
|
|
|||
|
|
|||
| class WeightedCrossEntropy(LanguageCrossEntropy): | |||
There was a problem hiding this comment.
I think a better approach for this would be to subclass torchmetrics.Metric directly since it would then be portable across Lightning, Composer and Torchtitan (at the very least), without needing composer installed. This is also pretty much overriding everything LanguageCrossEntropy does, so I see little advantage to it.
| @@ -12,6 +12,7 @@ dependencies: | |||
| - pytorch-cuda=12.1 | |||
| - pip: | |||
| - causal-conv1d==1.1.3.post1 | |||
| - mosaicml | |||
There was a problem hiding this comment.
Let's remove this re: https://github.com/Open-Athena/mdlm/pull/1/files#r1937418877. It's great that Composer is designed to be decoupled from models like that. Torchtitan is too AFAIK. Lightning is not, and I'm not sure what that means yet for running lightning models like this on other training frameworks yet. Either way, there shouldn't be any need to depend on composer in mdlm.
[WIP]