{"id":19450029,"url":"https://github.com/zipzou/hf-multitask-trainer","last_synced_at":"2025-10-26T12:05:29.743Z","repository":{"id":247804283,"uuid":"826859599","full_name":"zipzou/hf-multitask-trainer","owner":"zipzou","description":"The trainer for HF to record losses of different tasks and objectives.","archived":false,"fork":false,"pushed_at":"2025-03-12T07:14:02.000Z","size":41,"stargazers_count":46,"open_issues_count":3,"forks_count":3,"subscribers_count":1,"default_branch":"main","last_synced_at":"2025-10-11T17:23:16.845Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/zipzou.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2024-07-10T14:18:40.000Z","updated_at":"2025-09-25T04:31:04.000Z","dependencies_parsed_at":"2024-11-10T16:36:28.792Z","dependency_job_id":"bce3b1b2-a93d-4b68-91cd-247bcc0e66cc","html_url":"https://github.com/zipzou/hf-multitask-trainer","commit_stats":null,"previous_names":["zipzou/hf-multitask-trainer"],"tags_count":7,"template":false,"template_full_name":null,"purl":"pkg:github/zipzou/hf-multitask-trainer","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/zipzou%2Fhf-multitask-trainer","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/zipzou%2Fhf-multitask-trainer/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/zipzou%2Fhf-multitask-trainer/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/zipzou%2Fhf-multitask-trainer/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/zipzou","download_url":"https://codeload.github.com/zipzou/hf-multitask-trainer/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/zipzou%2Fhf-multitask-trainer/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":279008116,"owners_count":26084396,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-10-11T02:00:06.511Z","response_time":55,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-11-10T16:34:52.032Z","updated_at":"2025-10-11T17:23:18.346Z","avatar_url":"https://github.com/zipzou.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Enhanced Multitask Trainer for Separately Reporting Task's Metrics or Losses in HuggingFace Transformers\n\nThe [HuggingFace transformers](https://github.com/huggingface/transformers) library is widely used for model training. For example, to adapt a pretrained BERT model to a specific task domain, we often continue pretraining the model with two tasks in BERT: 1) Next Sentence Prediction (NSP) using the `[CLS]` token, and 2) Masked Language Modeling (MLM) using masked tokens. \n\n**A key issue is that the default `Trainer` in `transformers` assumes the first element of the output is the final loss to minimize. The loss returned by the `forward` method must be a scalar, so when training a multitask model like BERT, the loss needs to be combined.**\n\n**The `Trainer` class offers command-line arguments to control the training process. However, it only provides a combined loss value for all tasks, which obscures the individual losses of each task. This makes it challenging to monitor training and debug different task settings. Additionally, the `Tensorboard` report only shows the combined loss in its metrics.**\n\nTo facilitate multitask model training and review the loss of each task, as well as other training metrics, this trainer implementation is simple and useful.\n\nThe trainer works like the original `Trainer` in the `transformers` library. You just need to call the `report_metrics(...)` method to report the metrics that are important to you.\n\nBy the way, another utility you might need is [parser-binding](https://github.com/zipzou/parser-binding), which builds argument parsers from dataclasses and reads the arguments from command line scripts.\n\n## Usage\n\nFollow these steps to use the `HfMultiTaskTrainer`:\n\n1. Install the trainer:\n\n    ```sh\n    pip install hf-mtask-trainer\n    ```\n\n2. Replace the default trainer with `HfMultiTaskTrainer`:\n\n    ```python\n    from hf_mtask_trainer import HfMultiTaskTrainer\n\n    class Trainer(HfMultiTaskTrainer):\n        def __init__(...):\n            super().__init__(...)\n            # Additional initialization code\n    ```\n\n    Alternatively, you can directly instantiate the `HfMultiTaskTrainer`:\n\n    ```python\n    trainer = HfMultiTaskTrainer(...)\n    ```\n\n3. Report metrics in the model:\n\n    ```python\n    import torch.nn as nn\n\n    class Model(nn.Module):\n\n        supports_report_metrics: bool = True\n\n        def __init__(...):\n            super().__init__(...)\n            # Additional initialization code\n        \n        def forward(self, inputs, ...):\n            # Calculate metrics like loss, accuracy, etc.\n            task1_loss = ...\n            task2_loss = ...\n            acc = ...\n            f1 = ...\n            # Report the metrics\n            self.report_metrics(loss1=task1_loss, loss2=task2_loss, acc=acc, f1=f1)\n    ```\n\n    Add a flag `supports_report_metrics` where you need to report metrics, otherwise, the `report_metrics` would be not accessible.\n\n4. Start training the model:\n\n    As usual, call `trainer.train()` to start training.\n\nNow you can enjoy multitask training. If you set `--report tensorboard`, the metrics reported in the model will be displayed in Tensorboard diagrams.\n\n\n## Demo\n\nWe give a simple demo to mock a multi-task training in [test_trainer.py](./test_trainer.py).\n\nThe source code is:\n```python\nimport random\n\nimport numpy as np\nimport torch\nfrom torch import nn\nfrom torch.utils.data import Dataset\nfrom transformers.hf_argparser import HfArgumentParser\nfrom transformers.training_args import TrainingArguments\n\nfrom hf_mtask_trainer import HfMultiTaskTrainer\n\n# The model class\nclass TestModel(nn.Module):\n    supports_report_metrics: bool = True # IMPORTANT\n\n    def __init__(self, ) -\u003e None:\n        super().__init__()\n        self.scaler = nn.Parameter(torch.ones(1))\n\n    def forward(self, x):\n        test_tensor = x + self.scaler\n        test_np = np.array(np.random.randn()).astype(np.float32)\n        test_int = random.randint(1, 100)\n        test_float = random.random()\n        if hasattr(self, report_metrics): # checking if the report method is accessible or not is the robust practice\n            self.report_metrics(\n                tensor=test_tensor,\n                np=test_np,\n                integer=test_int,\n                fp_num=test_float\n            )\n\n        loss = ((\n            test_tensor + torch.from_numpy(test_np) + torch.tensor(test_int) +\n            torch.tensor(test_float) - 0\n        )).mean()\n\n        outputs = (loss, )\n\n        return outputs\n\n# Mock dataset\nclass MockDataset(Dataset):\n\n    def __len__(self):\n        return 1000\n\n    def __getitem__(self, index: int):\n        return dict(x=torch.randn(10, dtype=torch.float32))\n\n\ndef main():\n    parser = HfArgumentParser(TrainingArguments)\n    args, = parser.parse_args_into_dataclasses()\n    model = TestModel()\n    ds = MockDataset()\n    # Use HfMultiTaskTrainer rather than Trainer\n    trainer = HfMultiTaskTrainer(model, args, train_dataset=ds)\n\n    trainer.train()\n\n\nif __name__ == '__main__':\n    main()\n\n```\n\nRun the script to start training: `python test_trainer.py --output_dir ./test-output --per_device_train_batch_size 8 --gradient_accumulation_steps 4 --logging_steps 10 --num_train_epochs 10`.\n\nThe progress in the terminal is:\n```sh\n{'loss': 55.509, 'grad_norm': 1.0, 'learning_rate': 4.8387096774193554e-05, 'tensor': 0.9841784507036209, 'np': -0.21863683552946894, 'integer': 54.3, 'fp_num': 0.4434714410935673, 'epoch': 0.32}               \n{'loss': 48.2661, 'grad_norm': 1.0, 'learning_rate': 4.67741935483871e-05, 'tensor': 0.9985833063721656, 'np': -0.02904345905408263, 'integer': 46.725, 'fp_num': 0.5715085473120125, 'epoch': 0.64}              \n{'loss': 46.4612, 'grad_norm': 1.0, 'learning_rate': 4.516129032258064e-05, 'tensor': 0.9966234847903251, 'np': 0.010173140384722501, 'integer': 44.95, 'fp_num': 0.5043715258219943, 'epoch': 0.96}              \n{'loss': 48.6079, 'grad_norm': 1.0, 'learning_rate': 4.3548387096774194e-05, 'tensor': 0.9955430060625077, 'np': -0.03289987128227949, 'integer': 47.175, 'fp_num': 0.47028585139293366, 'epoch': 1.28}           \n{'loss': 50.091, 'grad_norm': 1.0, 'learning_rate': 4.1935483870967746e-05, 'tensor': 0.9734495922923088, 'np': 0.06655221048276871, 'integer': 48.55, 'fp_num': 0.5009696474466848, 'epoch': 1.6}                \n{'loss': 52.1638, 'grad_norm': 1.0, 'learning_rate': 4.032258064516129e-05, 'tensor': 1.0023577958345413, 'np': 0.18944044597446918, 'integer': 50.5, 'fp_num': 0.47205086657725437, 'epoch': 1.92}               \n{'loss': 61.3063, 'grad_norm': 1.0, 'learning_rate': 3.870967741935484e-05, 'tensor': 1.0168104887008667, 'np': -0.10900555825792253, 'integer': 60.0, 'fp_num': 0.39849607236524115, 'epoch': 2.24}              \n{'loss': 55.318, 'grad_norm': 1.0, 'learning_rate': 3.7096774193548386e-05, 'tensor': 1.015606315433979, 'np': 0.21950888196006418, 'integer': 53.575, 'fp_num': 0.5078790131376146, 'epoch': 2.56}               \n{'loss': 57.1703, 'grad_norm': 1.0, 'learning_rate': 3.548387096774194e-05, 'tensor': 1.0161942049860955, 'np': -0.08120755353011191, 'integer': 55.675, 'fp_num': 0.5603439507938002, 'epoch': 2.88}             \n{'loss': 47.6687, 'grad_norm': 1.0, 'learning_rate': 3.387096774193548e-05, 'tensor': 0.9780291050672532, 'np': 0.21060471932869404, 'integer': 46.025, 'fp_num': 0.4550899259063651, 'epoch': 3.2}               \n{'loss': 50.6742, 'grad_norm': 1.0, 'learning_rate': 3.2258064516129034e-05, 'tensor': 0.9773322150111199, 'np': 0.053728557180147615, 'integer': 49.15, 'fp_num': 0.4931880990797102, 'epoch': 3.52}             \n{'loss': 55.3104, 'grad_norm': 1.0, 'learning_rate': 3.0645161290322585e-05, 'tensor': 0.962137694656849, 'np': -0.079732296615839, 'integer': 53.975, 'fp_num': 0.45303205544101893, 'epoch': 3.84}              \n{'loss': 55.3539, 'grad_norm': 1.0, 'learning_rate': 2.9032258064516133e-05, 'tensor': 1.0214665666222573, 'np': -0.15776186664588748, 'integer': 53.9, 'fp_num': 0.590140296440284, 'epoch': 4.16}               \n{'loss': 49.332, 'grad_norm': 1.0, 'learning_rate': 2.7419354838709678e-05, 'tensor': 1.0191335454583168, 'np': -0.2712035422213376, 'integer': 48.025, 'fp_num': 0.5590896723075907, 'epoch': 4.48}              \n{'loss': 49.8865, 'grad_norm': 1.0, 'learning_rate': 2.5806451612903226e-05, 'tensor': 1.0170967370271682, 'np': 0.02669397685676813, 'integer': 48.275, 'fp_num': 0.5677363725430722, 'epoch': 4.8}              \n{'loss': 55.0644, 'grad_norm': 1.0, 'learning_rate': 2.4193548387096777e-05, 'tensor': 0.99910968542099, 'np': 0.12097712438553572, 'integer': 53.475, 'fp_num': 0.4693036682925622, 'epoch': 5.12}               \n{'loss': 56.9469, 'grad_norm': 1.0, 'learning_rate': 2.258064516129032e-05, 'tensor': 1.0159066557884215, 'np': -0.06122639870736748, 'integer': 55.6, 'fp_num': 0.3922143274213026, 'epoch': 5.44}               \n{'loss': 58.3238, 'grad_norm': 1.0, 'learning_rate': 2.0967741935483873e-05, 'tensor': 0.9946490600705147, 'np': -0.038768217992037536, 'integer': 56.875, 'fp_num': 0.49290766579450906, 'epoch': 5.76}          \n{'loss': 57.8349, 'grad_norm': 1.0, 'learning_rate': 1.935483870967742e-05, 'tensor': 0.9948656186461449, 'np': -0.15342782847583294, 'integer': 56.55, 'fp_num': 0.4434852700815277, 'epoch': 6.08}              \n{'loss': 57.5093, 'grad_norm': 1.0, 'learning_rate': 1.774193548387097e-05, 'tensor': 0.9814934283494949, 'np': 0.17727854922413827, 'integer': 55.85, 'fp_num': 0.5005189062297719, 'epoch': 6.4}                \n{'loss': 54.0808, 'grad_norm': 1.0, 'learning_rate': 1.6129032258064517e-05, 'tensor': 1.0003552585840225, 'np': 0.09905800204724073, 'integer': 52.425, 'fp_num': 0.5563636991813741, 'epoch': 6.72}             \n{'loss': 41.9312, 'grad_norm': 1.0, 'learning_rate': 1.4516129032258066e-05, 'tensor': 0.9884074732661248, 'np': 0.1483861011918634, 'integer': 40.275, 'fp_num': 0.51941084196083, 'epoch': 7.04}                \n{'loss': 54.1181, 'grad_norm': 1.0, 'learning_rate': 1.2903225806451613e-05, 'tensor': 1.0151973858475685, 'np': 0.47866107723675666, 'integer': 52.175, 'fp_num': 0.4492639089144623, 'epoch': 7.36}             \n{'loss': 50.6587, 'grad_norm': 1.0, 'learning_rate': 1.129032258064516e-05, 'tensor': 0.9820004492998123, 'np': -0.012274338398128748, 'integer': 49.2, 'fp_num': 0.4889366814531261, 'epoch': 7.68}              \n{'loss': 55.0801, 'grad_norm': 1.0, 'learning_rate': 9.67741935483871e-06, 'tensor': 0.9795809179544449, 'np': -0.07257360897492618, 'integer': 53.725, 'fp_num': 0.448120297698113, 'epoch': 8.0}                \n{'loss': 44.1352, 'grad_norm': 1.0, 'learning_rate': 8.064516129032258e-06, 'tensor': 0.9734664395451545, 'np': 0.286221909429878, 'integer': 42.375, 'fp_num': 0.5005484995389671, 'epoch': 8.32}                \n{'loss': 66.0453, 'grad_norm': 1.0, 'learning_rate': 6.451612903225806e-06, 'tensor': 0.9795126229524612, 'np': 0.030494442163035273, 'integer': 64.525, 'fp_num': 0.5103026267522799, 'epoch': 8.64}             \n{'loss': 56.4957, 'grad_norm': 1.0, 'learning_rate': 4.838709677419355e-06, 'tensor': 0.9856317490339279, 'np': 0.2455663602799177, 'integer': 54.775, 'fp_num': 0.48952095056963413, 'epoch': 8.96}              \n{'loss': 58.896, 'grad_norm': 1.0, 'learning_rate': 3.225806451612903e-06, 'tensor': 0.9927483782172203, 'np': 0.14382120433729143, 'integer': 57.275, 'fp_num': 0.4843773558505675, 'epoch': 9.28}               \n{'loss': 51.3854, 'grad_norm': 1.0, 'learning_rate': 1.6129032258064516e-06, 'tensor': 0.974456462264061, 'np': -0.03793883747421205, 'integer': 49.975, 'fp_num': 0.4738723365026173, 'epoch': 9.6}              \n{'loss': 51.1056, 'grad_norm': 1.0, 'learning_rate': 0.0, 'tensor': 1.0064959138631822, 'np': 0.17969400193542243, 'integer': 49.375, 'fp_num': 0.5444181010304485, 'epoch': 9.92}                                \n{'train_runtime': 0.4796, 'train_samples_per_second': 20852.484, 'train_steps_per_second': 646.427, 'train_loss': 53.31389662219632, 'epoch': 9.92}                                                               \n100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 310/310 [00:00\u003c00:00, 648.81it/s]\n```\n\n## Limitation\n\nThis trainer has not been fully tested yet but works for simple multitask training. Please report any issues if this plugin does not work for you.\n\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fzipzou%2Fhf-multitask-trainer","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fzipzou%2Fhf-multitask-trainer","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fzipzou%2Fhf-multitask-trainer/lists"}