[Feature request]: MultiTaskLitModule Refactor

[melo-gonzo]

Feature/behavior summary

As it stands, the MultiTaskLitModule can be difficult to work with and is fragile. A refactor could help to resolve some issues, as well as introduce more functionality, modularity, and overall quality of life improvements when working with the module. As this would be a more involved refactor, collecting pros/cons of the current setup and ideas for improvement in this issue would be ideal to get a bird's eye view of the task at hand.

Request attributes

• Would this be a refactor of existing code?
• Does this proposal require new package dependencies?
• Would this change break backwards compatibility?
• Does this proposal include a new model?
• Does this proposal include a new dataset?
• Does this proposal include a new task/workflow?

Related issues

No response

Solution description

The resolution to most of the discussion points may still be unknown, and this is primarily intended as a source of discussion, brainstorming, and idea collating.

(In no particular order)

• Pass only one encoder config, and generate output heads.
• Remove unused code (e.g. the forward pass, self.encoder attribute, etc.).
• Calculate loss and do backprop only if the optimizers associated dataset / task head had data processed. (see similar SAM Callback Update - Check If Optimizer Is Utilized #193)
• Unify or make clear the processing pipelines for multi vs. single data.
• Accept any arbitrary combination of dataset-task combinations

Additional notes

No response