week05_large_models/practice_part2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "bdfc71a449484eaeb3539fabb60a767c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_c5a57f756d474bdea5ce6b012edf1508",
              "IPY_MODEL_ce7ec364252f46b996ad0b32fe9d38cd",
              "IPY_MODEL_8d369b7163074ebd822d6702878256a4"
            ],
            "layout": "IPY_MODEL_ab8fc563343e4e7ea09356642de10756"
          }
        },
        "c5a57f756d474bdea5ce6b012edf1508": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_6a9ff0b43f794891ba1f5c4ecd6dcfdf",
            "placeholder": "​",
            "style": "IPY_MODEL_3787e2bdb8cd42a8bf1815a090ed9706",
            "value": "tokenizer_config.json: 100%"
          }
        },
        "ce7ec364252f46b996ad0b32fe9d38cd": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_b944e375f23d4ea98e58533f11857d33",
            "max": 50646,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_44d7e3d82cc54756a70a287bbd39ecb9",
            "value": 50646
          }
        },
        "8d369b7163074ebd822d6702878256a4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_6f4ac58e5676482e9bcb949adc4f84d1",
            "placeholder": "​",
            "style": "IPY_MODEL_d2f335cbff4e4ff2baa02a2ad7b82a3e",
            "value": " 50.6k/50.6k [00:00&lt;00:00, 2.96MB/s]"
          }
        },
        "ab8fc563343e4e7ea09356642de10756": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "6a9ff0b43f794891ba1f5c4ecd6dcfdf": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "3787e2bdb8cd42a8bf1815a090ed9706": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "b944e375f23d4ea98e58533f11857d33": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "44d7e3d82cc54756a70a287bbd39ecb9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "6f4ac58e5676482e9bcb949adc4f84d1": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d2f335cbff4e4ff2baa02a2ad7b82a3e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "1aa8ab0872514d57b0592ba43ea0105a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_bc7ba0334162415ba5f27f20deb28247",
              "IPY_MODEL_c0e6c92aaebc4f02a4943a572d8220dd",
              "IPY_MODEL_db12af0c8a0d41fa981dbb66e7ed49d4"
            ],
            "layout": "IPY_MODEL_4b7c453fb3994045a855b7370eee234f"
          }
        },
        "bc7ba0334162415ba5f27f20deb28247": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_c6e968957e8c4a21953a2d3b9ddd4f8d",
            "placeholder": "​",
            "style": "IPY_MODEL_2ac894dade0c48b7a30c058eded65b73",
            "value": "tokenizer.json: 100%"
          }
        },
        "c0e6c92aaebc4f02a4943a572d8220dd": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_94f4a642365b4e69aa15e7ea2bd1a29b",
            "max": 17209920,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_fa7d727b36b742419a928e271d403bba",
            "value": 17209920
          }
        },
        "db12af0c8a0d41fa981dbb66e7ed49d4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_6ee6f803e8a34061b52b09a1ed9c3a29",
            "placeholder": "​",
            "style": "IPY_MODEL_dcd4b9306a21422abdd2073913493326",
            "value": " 17.2M/17.2M [00:00&lt;00:00, 36.3MB/s]"
          }
        },
        "4b7c453fb3994045a855b7370eee234f": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "c6e968957e8c4a21953a2d3b9ddd4f8d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "2ac894dade0c48b7a30c058eded65b73": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "94f4a642365b4e69aa15e7ea2bd1a29b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "fa7d727b36b742419a928e271d403bba": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "6ee6f803e8a34061b52b09a1ed9c3a29": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "dcd4b9306a21422abdd2073913493326": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "d6509faf9220437e9371f83b5058fc29": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_d8e556d9c6c04603905217901cb38ac7",
              "IPY_MODEL_820a152271204ba18f6d4808e9fca90e",
              "IPY_MODEL_55e4bf33a30e400db72d47570e2345f8"
            ],
            "layout": "IPY_MODEL_5ca5324606e84da48635840a32e94f61"
          }
        },
        "d8e556d9c6c04603905217901cb38ac7": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_6a6840053fad4785a2124aea7c261f59",
            "placeholder": "​",
            "style": "IPY_MODEL_f1c7bf7153f14344a6e4a7a9bdc35b53",
            "value": "special_tokens_map.json: 100%"
          }
        },
        "820a152271204ba18f6d4808e9fca90e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_b26d8fd37493485a9c69a20de99262bf",
            "max": 459,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_76cbb91b2815478ea29d59b9e4bbc937",
            "value": 459
          }
        },
        "55e4bf33a30e400db72d47570e2345f8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_7b69fc37e8054da1977eebf3ab652157",
            "placeholder": "​",
            "style": "IPY_MODEL_0f03319b7e6f4fcc93e2b91ca3293584",
            "value": " 459/459 [00:00&lt;00:00, 14.2kB/s]"
          }
        },
        "5ca5324606e84da48635840a32e94f61": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "6a6840053fad4785a2124aea7c261f59": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "f1c7bf7153f14344a6e4a7a9bdc35b53": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "b26d8fd37493485a9c69a20de99262bf": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "76cbb91b2815478ea29d59b9e4bbc937": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "7b69fc37e8054da1977eebf3ab652157": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "0f03319b7e6f4fcc93e2b91ca3293584": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "8d891195bee042c5b3c88411367196ca": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_01dd79c9b6404d3c98affd6666ca2d22",
              "IPY_MODEL_50b1790060e94fab95d79317cc0dc333",
              "IPY_MODEL_f9649cfc2d54427ca3a70f87a82f4a5f"
            ],
            "layout": "IPY_MODEL_57f7eea22c934910b8897d340edf3107"
          }
        },
        "01dd79c9b6404d3c98affd6666ca2d22": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_fff708ec9f324f8980f8549639dae9ee",
            "placeholder": "​",
            "style": "IPY_MODEL_92c3644edc2c48aca666dc94b267dda8",
            "value": "config.json: 100%"
          }
        },
        "50b1790060e94fab95d79317cc0dc333": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_1c80c6a948d84ab2810ea45af8d27af2",
            "max": 935,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_74c395453ab64b0899505f152a12c50a",
            "value": 935
          }
        },
        "f9649cfc2d54427ca3a70f87a82f4a5f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_94a2a7b26e984778a2186a491bf06b4f",
            "placeholder": "​",
            "style": "IPY_MODEL_a7c591494d654cdd8a19141cf96d8036",
            "value": " 935/935 [00:00&lt;00:00, 56.1kB/s]"
          }
        },
        "57f7eea22c934910b8897d340edf3107": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "fff708ec9f324f8980f8549639dae9ee": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "92c3644edc2c48aca666dc94b267dda8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "1c80c6a948d84ab2810ea45af8d27af2": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "74c395453ab64b0899505f152a12c50a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "94a2a7b26e984778a2186a491bf06b4f": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "a7c591494d654cdd8a19141cf96d8036": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "d77be2d56f3846d4991c92ee953731ec": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_f940bf4b3d454ad7a3592b0cc51979fd",
              "IPY_MODEL_8339128e688d4221bc318b45910404a6",
              "IPY_MODEL_0309561fa3734a0ea3cd3421c420ff56"
            ],
            "layout": "IPY_MODEL_2329d7fe319b4383ab87aabe5e3e17fa"
          }
        },
        "f940bf4b3d454ad7a3592b0cc51979fd": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_df6158bafdd94f4e8b9f147acc0ef8f5",
            "placeholder": "​",
            "style": "IPY_MODEL_48669f5193124c7ca52d554a998be789",
            "value": "model.safetensors: 100%"
          }
        },
        "8339128e688d4221bc318b45910404a6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_a12be7a58f764dff96d41ba35d3a9c76",
            "max": 2471645608,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_ad6a69b6a4614050a497eccebdc920b2",
            "value": 2471645608
          }
        },
        "0309561fa3734a0ea3cd3421c420ff56": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_37919cdcd7304db2a60df1cc387ca4ce",
            "placeholder": "​",
            "style": "IPY_MODEL_d1cbd32f079042398389dd2f250cc8ef",
            "value": " 2.47G/2.47G [00:59&lt;00:00, 42.7MB/s]"
          }
        },
        "2329d7fe319b4383ab87aabe5e3e17fa": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "df6158bafdd94f4e8b9f147acc0ef8f5": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "48669f5193124c7ca52d554a998be789": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "a12be7a58f764dff96d41ba35d3a9c76": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "ad6a69b6a4614050a497eccebdc920b2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "37919cdcd7304db2a60df1cc387ca4ce": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d1cbd32f079042398389dd2f250cc8ef": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "0b353dfbb9784b3883ff31ed7124f450": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_443a315962ab46ea8b6d8f2ed354038e",
              "IPY_MODEL_d1d68bb4f42647e39806f1fc4dd9c014",
              "IPY_MODEL_ede69839385f48399d51b0c9a0568353"
            ],
            "layout": "IPY_MODEL_5768b9b7ffec4884869b2334eb946015"
          }
        },
        "443a315962ab46ea8b6d8f2ed354038e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_9045df041cac425fa7e109d303bb4d2f",
            "placeholder": "​",
            "style": "IPY_MODEL_0dc89a0bd4e34ce1b60597cb1acd7fdf",
            "value": "generation_config.json: 100%"
          }
        },
        "d1d68bb4f42647e39806f1fc4dd9c014": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_5a5c1dcd5a214c93bb917c8eb9029af2",
            "max": 230,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_c5f9ad4f3dd94e8b8d05530f11fed714",
            "value": 230
          }
        },
        "ede69839385f48399d51b0c9a0568353": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_fd64196b94f64bf59d4bbd1307aaf05c",
            "placeholder": "​",
            "style": "IPY_MODEL_d643e91490274b1988b242e4d784ff82",
            "value": " 230/230 [00:00&lt;00:00, 15.4kB/s]"
          }
        },
        "5768b9b7ffec4884869b2334eb946015": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "9045df041cac425fa7e109d303bb4d2f": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "0dc89a0bd4e34ce1b60597cb1acd7fdf": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "5a5c1dcd5a214c93bb917c8eb9029af2": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "c5f9ad4f3dd94e8b8d05530f11fed714": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "fd64196b94f64bf59d4bbd1307aaf05c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d643e91490274b1988b242e4d784ff82": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mryab/efficient-dl-systems/blob/main/week05_large_models/practice_part2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Efficient DL Practice: Advanced Parallelism (5 points)\n",
        "\n",
        "In this practice session, we'll cover techniques for training large models in parallel: **Model** and **Sequence Parallelism**.\n",
        "More precisely, you will implement them, and we will root for you as you go. Good luck, 🥩👜!\n",
        "\n"
      ],
      "metadata": {
        "id": "RlQVQby9YJip"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# dependencies: the code will likely work with slightly newer/older versions, but may require minimal patching\n",
        "%pip install -q transformers==4.48.3 peft==0.14.0\n",
        "\n",
        "import transformers; assert transformers.__version__.startswith(\"4.48\"), transformers.__version__\n",
        "import peft; assert peft.__version__.startswith(\"0.14\"), peft.__version__"
      ],
      "metadata": {
        "id": "9cp865qOmgcl",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "02202d9e-12d4-4341-b31f-006c356bd2b7"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[33mWARNING: Ignoring invalid distribution ~vidia-cusparse-cu12 (/usr/local/lib/python3.11/dist-packages)\u001b[0m\u001b[33m\n",
            "\u001b[0m\u001b[33mWARNING: Ignoring invalid distribution ~vidia-cusparse-cu12 (/usr/local/lib/python3.11/dist-packages)\u001b[0m\u001b[33m\n",
            "\u001b[0m\u001b[33mWARNING: Ignoring invalid distribution ~vidia-cusparse-cu12 (/usr/local/lib/python3.11/dist-packages)\u001b[0m\u001b[33m\n",
            "\u001b[0m"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "__Part 1: Tensor Parallelism (2 points)__\n",
        "![img](https://pytorch.org/tutorials/_images/megatron_lm.png)\n",
        "\n",
        "We'll begin by implementing a simple tensor parallelism (also known as the [original](https://papers.nips.cc/paper_files/paper/2012/hash/c399862d3b9d6b76c8436e924a68c45b-Abstract.html) model parallelism).\n",
        "\n",
        "Our ultimate objective is to run and fine-tune a Llama 3.x model in tensor-parallel mode. However, it is rather difficult to do that in one go, especially if you take bugs into account. So we'll start simple: __here's a single Llama MLP module:__\n",
        "\n",
        "`please read the code below carefully, it's a template for the remaining assgnments`."
      ],
      "metadata": {
        "id": "tw_-8huWabX1"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KbcTjrsRB-ps",
        "outputId": "48576f02-d041-42be-dc44-104e8d84251a"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Overwriting tensor_parallel_mlp.py\n"
          ]
        }
      ],
      "source": [
        "%%writefile tensor_parallel_mlp.py\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.distributed as dist\n",
        "\n",
        "\n",
        "class LlamaMLP(nn.Module):  #  based on llama 3.1 8B configuration\n",
        "    def __init__(self, hidden_size: int = 4096, intermediate_size: int = 14336):\n",
        "        super().__init__()\n",
        "        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)\n",
        "        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)\n",
        "        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)\n",
        "\n",
        "    def forward(self, input):\n",
        "        return self.down_proj(F.silu(self.gate_proj(input)) * self.up_proj(input))\n",
        "\n",
        "\n",
        "class ComputeWithAllReduce(torch.autograd.Function):\n",
        "    @staticmethod  # fun fact: torch.distributed.nn has differentiable all_reduce!\n",
        "    def forward(ctx, tp_shard: nn.Module, input: torch.Tensor):\n",
        "        input = input.detach().requires_grad_(input.requires_grad)\n",
        "        ctx.save_for_backward(input)\n",
        "        ctx._tp_shard = tp_shard\n",
        "        output = tp_shard(input)\n",
        "        dist.all_reduce(output)\n",
        "        return output\n",
        "    @staticmethod\n",
        "    def backward(ctx, grad_output: torch.Tensor):\n",
        "        with torch.enable_grad():\n",
        "          output = ctx._tp_shard(ctx.saved_tensors[0])\n",
        "          output.backward(grad_output)\n",
        "        dist.all_reduce(ctx.saved_tensors[0].grad)\n",
        "        return None, ctx.saved_tensors[0].grad\n",
        "\n",
        "\n",
        "class AllReduceModule(nn.Sequential):\n",
        "    def forward(self, input: torch.Tensor):\n",
        "        return ComputeWithAllReduce.apply(super().forward, input)\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    dist.init_process_group(\"gloo\")   # use nccl for cuda devices\n",
        "    torch.manual_seed(1337)           # init weights equally on all ranks\n",
        "    rank, world_size = dist.get_rank(), dist.get_world_size()\n",
        "\n",
        "    for active_rank in range(world_size):\n",
        "      dist.barrier()  # initialize each rank sequentially to save system RAM\n",
        "      if rank != active_rank: continue\n",
        "\n",
        "      # we will now implement Tensor Parallelism for the ref_module below:\n",
        "      ref_module = nn.Sequential(nn.RMSNorm(4096), LlamaMLP())\n",
        "      # compute reference tensors to test against them later\n",
        "      input = torch.randn(1, 4096, requires_grad=True)\n",
        "      ref_output = ref_module(input)\n",
        "      ref_output.sum().backward()\n",
        "      ref_input_grad = input.grad.clone()\n",
        "\n",
        "      # TP step 1: define a module that computes a portion of intermediate units\n",
        "      intermediate_size = ref_module[1].down_proj.in_features\n",
        "      local_units = intermediate_size // world_size\n",
        "      assert intermediate_size % world_size == 0\n",
        "      tp_module = nn.Sequential(   # assign a portion of units per rank --v\n",
        "          nn.RMSNorm(4096), AllReduceModule(LlamaMLP(intermediate_size=local_units))\n",
        "      )   # all-reduce outputs during forward, all-reduce gradients on backward\n",
        "\n",
        "      with torch.no_grad():  # copy select weights from the reference MLP\n",
        "        # v-- input norm layer is too small to bother parallelizing - we replicate it!\n",
        "        tp_module[0].load_state_dict(ref_module[0].state_dict())\n",
        "        # up and gate projections are sharded across output units\n",
        "        unit_slice = slice(local_units * rank, local_units * (rank + 1))\n",
        "        tp_module[1][0].up_proj.weight[...] = ref_module[1].up_proj.weight[unit_slice]\n",
        "        tp_module[1][0].gate_proj.weight[...] = ref_module[1].gate_proj.weight[unit_slice]\n",
        "        # down projection is sharded across input units, matching up/gate proj\n",
        "        tp_module[1][0].down_proj.weight[...] = ref_module[1].down_proj.weight[:, unit_slice]\n",
        "      print(f\"Initialized {rank=}\", flush=True)\n",
        "      del ref_module  # free RAM for next rank\n",
        "\n",
        "    dist.barrier()  # test 1: forward pass\n",
        "    tp_input = input.detach().requires_grad_(True)\n",
        "    tp_output = tp_module(tp_input)\n",
        "    if rank == 0:\n",
        "        print(f\"\\nReference outputs ({rank=}):\", ref_output.data, flush=True)\n",
        "    for i in range(world_size):\n",
        "        dist.barrier()\n",
        "        if i != rank: continue\n",
        "        print(f\"TParallel outputs ({rank=}):\", tp_output.data, flush=True)\n",
        "        assert torch.allclose(tp_output, ref_output, atol=1e-6), f\"output mismatch on {rank=}\"\n",
        "\n",
        "    dist.barrier()  # test 2: backward w.r.t. inputs\n",
        "    assert tp_input.grad is None\n",
        "    tp_output.sum().backward()\n",
        "    if rank == 0:\n",
        "        print(f\"\\nReference input grad ({rank=}):\", ref_input_grad, flush=True)\n",
        "    for i in range(world_size):\n",
        "        dist.barrier()\n",
        "        if i != rank: continue\n",
        "        print(f\"TParallel input grad ({rank=}):\", tp_input.grad.data, flush=True)\n",
        "        assert torch.allclose(tp_input.grad, ref_input_grad, atol=1e-6), f\"input_grad mismatch on {rank=}\"\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!OMP_NUM_THREADS=1 torchrun --nproc_per_node 4 tensor_parallel_mlp.py"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4z99fF4HC9hp",
        "outputId": "e4d08391-746c-47a5-f137-c4b1e56b64c9"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Initialized rank=0\n",
            "Initialized rank=1\n",
            "Initialized rank=2\n",
            "Initialized rank=3\n",
            "\n",
            "Reference outputs (rank=0): tensor([[-0.1145,  0.0160,  0.0500,  ..., -0.1455,  0.1126, -0.0192]])\n",
            "TParallel outputs (rank=0): tensor([[-0.1145,  0.0160,  0.0500,  ..., -0.1455,  0.1126, -0.0192]])\n",
            "TParallel outputs (rank=1): tensor([[-0.1145,  0.0160,  0.0500,  ..., -0.1455,  0.1126, -0.0192]])\n",
            "TParallel outputs (rank=2): tensor([[-0.1145,  0.0160,  0.0500,  ..., -0.1455,  0.1126, -0.0192]])\n",
            "TParallel outputs (rank=3): tensor([[-0.1145,  0.0160,  0.0500,  ..., -0.1455,  0.1126, -0.0192]])\n",
            "\n",
            "Reference input grad (rank=0): tensor([[ 0.0343, -0.2492, -0.1858,  ..., -0.0541,  0.0388, -0.1529]])\n",
            "TParallel input grad (rank=0): tensor([[ 0.0343, -0.2492, -0.1858,  ..., -0.0541,  0.0388, -0.1529]])\n",
            "TParallel input grad (rank=1): tensor([[ 0.0343, -0.2492, -0.1858,  ..., -0.0541,  0.0388, -0.1529]])\n",
            "TParallel input grad (rank=2): tensor([[ 0.0343, -0.2492, -0.1858,  ..., -0.0541,  0.0388, -0.1529]])\n",
            "TParallel input grad (rank=3): tensor([[ 0.0343, -0.2492, -0.1858,  ..., -0.0541,  0.0388, -0.1529]])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Note that the code above lacks two details:\n",
        "- it uses a form of checkpointing, but does not save random state, which would be required if you use dropout;\n",
        "- it replicates RMSNorm, but it is not synchronized. Training would require all-reduce-ing gradients for those layers, e.g. by wrapping them with DDP.\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "\n",
        "__Task 1 (1 point):__ Implement tensor-parallel multi-head attention.\n",
        "\n",
        "Like with the MLP module before, you can partition attention across multiple devices. This time, every device is to compute a portion of whole attention **heads** (and not individual units). We exploit the property that an multi-head attention layer can be viewed as a sum of individual head outputs after output projection.\n",
        "\n",
        "For the sake of formality, this is the computation you need to parallelize:"
      ],
      "metadata": {
        "id": "nMnNI9I5d1R_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from transformers.models.llama.modeling_llama import LlamaConfig, LlamaAttention, LlamaRotaryEmbedding\n",
        "MODEL_NAME = \"unsloth/Llama-3.2-1B\"  # for testing (but not grading!), you may want to use Maykeye/TinyLLama-v0\n",
        "config = LlamaConfig.from_pretrained(MODEL_NAME)\n",
        "layer = LlamaAttention(config, layer_idx=5)\n",
        "rotary_emb = LlamaRotaryEmbedding(config)\n",
        "\n",
        "input = torch.randn(1, 128, config.hidden_size, requires_grad=True)\n",
        "position_embeddings = rotary_emb(input, position_ids=torch.arange(128)[None])\n",
        "\n",
        "output, *_etc = layer(input, attention_mask=None, position_embeddings=position_embeddings)\n",
        "print(f\"{output=}\")\n",
        "output.norm().backward()\n",
        "print(f\"{input.grad=}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ub-VLG38kpsl",
        "outputId": "e488fc79-0168-4d42-dd41-8f7892de2958"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.11/dist-packages/huggingface_hub/utils/_auth.py:94: UserWarning: \n",
            "The secret `HF_TOKEN` does not exist in your Colab secrets.\n",
            "To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session.\n",
            "You will be able to reuse this secret in all of your notebooks.\n",
            "Please note that authentication is recommended but still optional to access public models or datasets.\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "output=tensor([[[ 0.0096, -0.0238,  0.0242,  ..., -0.0217, -0.0136,  0.0237],\n",
            "         [-0.0025, -0.0019,  0.0408,  ..., -0.0186,  0.0170,  0.0290],\n",
            "         [ 0.0110, -0.0258,  0.0225,  ..., -0.0073,  0.0033,  0.0340],\n",
            "         ...,\n",
            "         [ 0.0148, -0.0118,  0.0502,  ..., -0.0153, -0.0119,  0.0338],\n",
            "         [ 0.0071, -0.0040,  0.0366,  ..., -0.0214, -0.0104,  0.0600],\n",
            "         [-0.0003, -0.0279,  0.0330,  ..., -0.0091, -0.0124,  0.0251]]],\n",
            "       grad_fn=<UnsafeViewBackward0>)\n",
            "input.grad=tensor([[[ 0.0020,  0.0006, -0.0003,  ..., -0.0004, -0.0011, -0.0018],\n",
            "         [ 0.0020,  0.0005, -0.0002,  ..., -0.0005, -0.0011, -0.0016],\n",
            "         [ 0.0020,  0.0006, -0.0002,  ..., -0.0005, -0.0012, -0.0017],\n",
            "         ...,\n",
            "         [ 0.0019,  0.0006, -0.0002,  ..., -0.0004, -0.0011, -0.0018],\n",
            "         [ 0.0019,  0.0007, -0.0002,  ..., -0.0004, -0.0012, -0.0017],\n",
            "         [ 0.0020,  0.0007, -0.0002,  ..., -0.0004, -0.0011, -0.0018]]])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Same as before, your task is to create a multi-head attention layer, partition it across ranks and verify two things:\n",
        "- attention outputs on the same inputs (and mask) match with the non-parallel version;\n",
        "- gradients w.r.t. attention inputs are the same; gradients w.r.t. mask need not be verified.\n"
      ],
      "metadata": {
        "id": "uRYcD0qxmTuQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%writefile tensor_parallel_attn.py\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.distributed as dist\n",
        "\n",
        "\n",
        "class MyLlamaAttention(nn.Module):\n",
        "    ...  # please take a reference implementation of Llama attention from Hugging Face transformers:\n",
        "    # https://github.com/huggingface/transformers/blob/v4.44-release/src/transformers/models/llama/modeling_llama.py#L326-L455\n",
        "    # You can also directly import transformers.models.llama.modeling_llama.LlamaAttention, as in the reference above.\n",
        "    # Alternatively, you are welcome to simplify their code or implement your own version.\n",
        "\n",
        "    # Note: the link above points to an older version of attention with built-in rotary position embeddings (RoPE);\n",
        "    # If you are using a newer version, please make sure to define extra inputs\n",
        "\n",
        "\n",
        "# You will likely need to define additional classes below, e.g. a module to perform all-reduce\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    dist.init_process_group(\"gloo\")   # use nccl for cuda devices\n",
        "    torch.manual_seed(1337)           # init weights equally on all ranks\n",
        "    rank, world_size = dist.get_rank(), dist.get_world_size()\n",
        "\n",
        "    for active_rank in range(world_size):\n",
        "      dist.barrier()  # initialize each rank sequentially to save system RAM\n",
        "      if rank != active_rank: continue\n",
        "\n",
        "      # we will now implement Tensor Parallelism for the ref_module below:\n",
        "      ref_module = MyLlamaAttention()\n",
        "      # ^-- you may need to modify this code, e.g. pass parameters or use transformers LlamaAttention (as above)\n",
        "\n",
        "      # generate reference tensors to test against them later\n",
        "      input = torch.randn(1, 128, 4096, requires_grad=True)\n",
        "      extra_inputs = dict()  # <-- OPTIONAL: either design additional inputs here, as in the reference above\n",
        "\n",
        "      ref_output = ref_module(input, **extra_inputs)\n",
        "      ref_output.sum().backward()\n",
        "      ref_input_grad = input.grad.clone()\n",
        "\n",
        "      # TP step 1: define a module that computes a portion of attention heads\n",
        "\n",
        "      tp_module = <YOUR CODE HERE>  # create a tensor-parallel version of the Attention module\n",
        "\n",
        "      with torch.no_grad():\n",
        "          <YOUR CODE HERE>  # copy select weights from the reference attention\n",
        "\n",
        "      print(f\"Initialized {rank=}\", flush=True)\n",
        "      del ref_module  # free RAM for next rank\n",
        "\n",
        "    # TEST AREA: you are free to add additional parameters, but your code *must* run the same tests as below\n",
        "    dist.barrier()  # test 1: forward pass\n",
        "    tp_input = input.detach().requires_grad_(True)\n",
        "    tp_output = tp_module(tp_input, **extra_inputs)\n",
        "    if rank == 0:\n",
        "        print(f\"\\nReference outputs ({rank=}):\", ref_output.data, flush=True)\n",
        "    for i in range(world_size):\n",
        "        dist.barrier()\n",
        "        if i != rank: continue\n",
        "        print(f\"TParallel outputs ({rank=}):\", tp_output.data, flush=True)\n",
        "        assert torch.allclose(tp_output, ref_output, atol=1e-5), f\"output mismatch on {rank=}\"\n",
        "\n",
        "    dist.barrier()  # test 2: backward w.r.t. inputs\n",
        "    assert tp_input.grad is None\n",
        "    tp_output.sum().backward()\n",
        "    if rank == 0:\n",
        "        print(f\"\\nReference input grad ({rank=}):\", ref_input_grad, flush=True)\n",
        "    for i in range(world_size):\n",
        "        dist.barrier()\n",
        "        if i != rank: continue\n",
        "        print(f\"TParallel input grad ({rank=}):\", tp_input.grad.data, flush=True)\n",
        "        assert torch.allclose(tp_input.grad, ref_input_grad, atol=1e-4), f\"input_grad mismatch on {rank=}\"\n"
      ],
      "metadata": {
        "id": "6SrcK3avd3GA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!OMP_NUM_THREADS=1 torchrun --nproc_per_node 2 tensor_parallel_attn.py\n",
        "# ^-- feel free to modify parameters, as long as there are at least 2 ranks"
      ],
      "metadata": {
        "id": "yodlDa_cnPWH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Well done! *(hopefully. If not, go back and, well... do it)*\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "\n",
        "### Full model conversion\n",
        "\n",
        "Now let's apply this technique to parallelize the actual Llama model. As in, with weights.\n",
        "\n",
        "__Task 2 (1 point):__ Combine the two previous techniques in one file that parallelizes an actual Llama model and .generates meaningful output. For simplicity, you do not need to partition key-value cache here - only the forward pass itself. We will default to generating tokens with recomputation.\n",
        "\n",
        "For the sake of formality, your task is to parallelize the following inference code:\n"
      ],
      "metadata": {
        "id": "lxTbiVvInO9i"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import transformers\n",
        "MODEL_NAME = \"unsloth/Llama-3.2-1B\"  # for testing (but not grading!), you may want to use Maykeye/TinyLLama-v0\n",
        "\n",
        "tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME)\n",
        "model = transformers.LlamaForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float32)  # <-- you are allowed to switch to bf16\n",
        "\n",
        "prompt = \"A quick brown fox\"\n",
        "input_ids = tokenizer(prompt, return_tensors='pt')[\"input_ids\"]\n",
        "print(end=prompt)\n",
        "for i in range(5):\n",
        "  with torch.no_grad():\n",
        "    new_token = model(input_ids).logits[0, -1].argmax(-1)\n",
        "    input_ids = torch.cat([input_ids, new_token.view(1, 1)], dim=1)\n",
        "  print(end=tokenizer.decode(new_token), flush=True)\n",
        "# pro tip: delete the model or restart session to free RAM for the TP experiments"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "152wIY8fo0TI",
        "outputId": "09394465-9b0e-4927-eb4c-343d0b61278e"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "A quick brown fox jumps over the lazy dog"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "**Requirements:** your code must do the following things for the full grade:\n",
        "- instantiate an actually trained Llama model (Llama 3.2 1B or larger is fine, maykeye is not)\n",
        "- run forward pass with at least 2 ranks and verify that the logits are close,\n",
        "- run backward pass w.r.t. non-parallelized input embeddings, verify that the gradients are close,\n",
        "- perform inference for 10 steps to verify that the model produces meaningful outputs (see below)\n",
        "\n",
        "You are only required to tensor-parallel-ize the transformer layers. Parallelizing embeddings and logits is optional. If you do choose to parallelize embeddings, we sincerely recommend that you partition across the embedding dim, not across tokens - so that the computation is balanced."
      ],
      "metadata": {
        "id": "NLtaTR6zo6VU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "<<A whole lot of your code here>>"
      ],
      "metadata": {
        "id": "2cBtEb4Bd3Ox"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "<<... and a dedicated cell to show off that it works>>"
      ],
      "metadata": {
        "id": "plldOGaGo8A-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "### Using [`torch.distributed.tensor`](https://pytorch.org/docs/stable/distributed.tensor.html)\n",
        "\n",
        "PyTorch has an in-built functionality called [DTensor](https://pytorch.org/docs/stable/distributed.tensor.html), designed to help implementing tensor-level parallelism with various sharding strategies. This includes Tensor parallelism itself, as well as other techniques such as Sequence Parallelism, as they are both, essentially, parallelism across different tensor dimensions.\n",
        "\n",
        "__Task 3 (1 point):__ Your next task will be to replicate your previous code (llama inference) using DTensor instead of manual AllReduce. We recommend you start by skimming the [documentation for DTensor](https://pytorch.org/docs/stable/distributed.tensor.html) to learn the interface and [the minimal example](https://github.com/pytorch/examples/blob/main/distributed/tensor_parallelism/tensor_parallel_example.py) to learn how to put the pieces together.\n",
        "\n",
        "\n",
        "We recommend that you dedicate some time to learn and play with it before you proceed to parallelize Llama.\n",
        "\n",
        "The main objective is the same as in the previous task - run .generate with DTensor - and then compare it against the manual implementation. **Please report at least some speed comparison for forward and backward passes between this and the previous task.** If absolutely impossible (e.g. you don't have multiple gpus), we can accept a fallback assignment of implementing basic training: overfit the model to a single batch (like task 5 below) and demonstrate that it works - if you choose this option, say so in bold, large-font letters somewhere where the grader can see.\n",
        "\n",
        "But first, here's a quick demo of using DTensor for simple matrix multiplication - meant as a testbed for your experiments."
      ],
      "metadata": {
        "id": "hxBHNBsDq2r2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%writefile tensor_parallel_mlp_dtensor.py\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.distributed as dist\n",
        "from torch.distributed.device_mesh import init_device_mesh\n",
        "from torch.distributed.tensor import DTensor, DeviceMesh, Replicate, Shard\n",
        "import torch.distributed.tensor.parallel as tp\n",
        "\n",
        "\n",
        "class LlamaMLP(nn.Module):  # same module, but with smaller dims for quick prototyping\n",
        "    def __init__(self, hidden_size: int = 1024, intermediate_size: int = 4096):\n",
        "        super().__init__()\n",
        "        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)\n",
        "        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)\n",
        "        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)\n",
        "\n",
        "    def forward(self, input):\n",
        "        return self.down_proj(F.silu(self.gate_proj(input)) * self.up_proj(input))\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    dist.init_process_group(\"gloo\")  # use nccl for cuda devices\n",
        "    torch.manual_seed(1337)          # init weights equally on all ranks\n",
        "    rank, world_size = dist.get_rank(), dist.get_world_size()\n",
        "\n",
        "    # Initialize device mesh for tensor parallelism\n",
        "    device_mesh = init_device_mesh(device_type=\"cpu\", mesh_shape=(world_size,))  # use \"cuda\" for GPU\n",
        "\n",
        "    # Create reference module for comparison\n",
        "    ref_module = nn.Sequential(nn.RMSNorm(1024), LlamaMLP())\n",
        "\n",
        "    input = torch.randn(1, 1024, requires_grad=True)\n",
        "    ref_output = ref_module(input)\n",
        "    ref_output.sum().backward()\n",
        "    ref_input_grad = input.grad.clone()\n",
        "\n",
        "    # Create tensor parallel module (we wrap ref_module instead of copying)\n",
        "    tp_module = tp.parallelize_module(\n",
        "        ref_module,\n",
        "        device_mesh,\n",
        "        parallelize_plan={  # define parallelism type for each module\n",
        "            # up_proj and gate_proj are column-wise parallel (sharded across outputs);\n",
        "            \"1.up_proj\": tp.ColwiseParallel(),\n",
        "            \"1.gate_proj\": tp.ColwiseParallel(),\n",
        "            # down_proj is row-wise parallel (sharded across input dim)\n",
        "            \"1.down_proj\": tp.RowwiseParallel(),\n",
        "          },  # note: RMSNorm is simply replicated across all devices - hence, we skip it\n",
        "    )\n",
        "    if rank == 0:  # Note: no need to copy weight chunks manually: DTensor handles parameter sharding for us\n",
        "      for name, param in tp_module.named_parameters():\n",
        "        print(f\"{name=},\\ttype={type(param.data)}\\tglobal shape={param.shape},\\tlocal shape={param._local_tensor.shape if hasattr(param, '_local_tensor') else param.shape}\")\n",
        "\n",
        "    dist.barrier()  # Test forward and backward pass with Tensor Parallelism\n",
        "    tp_input = input.detach().requires_grad_(True)\n",
        "    tp_output = tp_module(tp_input)\n",
        "    tp_output.sum().backward()\n",
        "    tp_output = tp_output.trigger_wait()  # convert from AsyncCollectiveTensor to regular torch tensor\n",
        "    if rank == 0:\n",
        "        print(f\"\\nReference outputs ({rank=}):\", ref_output.data, flush=True)\n",
        "        print(f\"TParallel outputs ({rank=}):\", tp_output.data, flush=True)\n",
        "        print(f\"\\nReference input grad ({rank=}):\", ref_input_grad, flush=True)\n",
        "        print(f\"TParallel input grad ({rank=}):\", tp_input.grad, flush=True)\n",
        "    dist.barrier()\n",
        "    assert torch.allclose(tp_output, ref_output, atol=1e-6), f\"output mismatch on {rank=}\"\n",
        "    assert torch.allclose(tp_input.grad, ref_input_grad, atol=1e-6), f\"input_grad mismatch on {rank=}\"\n",
        "    print(end=f\"Tests passed ({rank=})\\n\", flush=True); dist.barrier()\n",
        "\n",
        "# fun fact: 90% of the code above was generated by grok-3 for prompt \"Please rewrite the following code using torch.distributed.tensor ```python <paste MLP code here>```\"\n",
        "# the remaining 10% are nasty bugfixes that took 99% of assignment preparation time. Do not trust the shogoths yet :)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CtvFF26mDWH2",
        "outputId": "9ddddac6-65aa-4341-f12f-e884a26f0bea"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Overwriting tensor_parallel_mlp_dtensor.py\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!OMP_NUM_THREADS=1 torchrun --nproc_per_node 2 tensor_parallel_mlp_dtensor.py"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rbGsyCMHDtXY",
        "outputId": "337e86a1-d9fa-48e2-aff0-2f0089b56ff2"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/usr/local/lib/python3.11/dist-packages/torch/distributed/tensor/_random.py:45: UserWarning: DTensor random operators may not have complete support on cpu device mesh\n",
            "  warnings.warn(\n",
            "/usr/local/lib/python3.11/dist-packages/torch/distributed/tensor/_random.py:45: UserWarning: DTensor random operators may not have complete support on cpu device mesh\n",
            "  warnings.warn(\n",
            "name='0.weight',\ttype=<class 'torch.Tensor'>\tglobal shape=torch.Size([1024]),\tlocal shape=torch.Size([1024])\n",
            "name='1.gate_proj.weight',\ttype=<class 'torch.distributed.tensor.DTensor'>\tglobal shape=torch.Size([4096, 1024]),\tlocal shape=torch.Size([2048, 1024])\n",
            "name='1.up_proj.weight',\ttype=<class 'torch.distributed.tensor.DTensor'>\tglobal shape=torch.Size([4096, 1024]),\tlocal shape=torch.Size([2048, 1024])\n",
            "name='1.down_proj.weight',\ttype=<class 'torch.distributed.tensor.DTensor'>\tglobal shape=torch.Size([1024, 4096]),\tlocal shape=torch.Size([1024, 2048])\n",
            "\n",
            "Reference outputs (rank=0): tensor([[ 0.0102,  0.0432, -0.0467,  ...,  0.0798, -0.0179,  0.0527]])\n",
            "TParallel outputs (rank=0): tensor([[ 0.0102,  0.0432, -0.0467,  ...,  0.0798, -0.0179,  0.0527]])\n",
            "\n",
            "Reference input grad (rank=0): tensor([[ 0.1543, -0.0858, -0.0882,  ..., -0.2082,  0.0298,  0.2388]])\n",
            "TParallel input grad (rank=0): tensor([[ 0.1543, -0.0858, -0.0882,  ..., -0.2082,  0.0298,  0.2388]])\n",
            "Tests passed (rank=0)\n",
            "Tests passed (rank=1)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "<<A WHOLE LOT OF YOUR CODE - apply the same strategy to implement tensor-only parallelism for full LLama model>>"
      ],
      "metadata": {
        "id": "qoA-ZAW4doMy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "<<YOUR CODE - run the script, check correctness, then run text generation with tensor-parallel code>>"
      ],
      "metadata": {
        "id": "DYPmr0xCfgEK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "\n",
        "\n",
        "### Sequence Parallelism with Ulysses\n",
        "\n",
        "\n",
        "Now let's parallelize the other way - across the sequence dimension. To showcase why this is necessary, our main task will be to parallelize LLM fine-tuning over a very long sequence. The way you do this, of course, is through Sequence Parallelism. You can implement naive [sequence parallelism](https://arxiv.org/abs/2205.05198), similar to [DeepSpeed Ulysses](https://arxiv.org/pdf/2309.14509) (n.b.: not the first work to do this).\n",
        "\n",
        "![figure-from-paper](https://ar5iv.labs.arxiv.org/html/2309.14509/assets/figs/image3.png)\n",
        "\n",
        "\n",
        "Here's the short version:\n",
        "- All weights are replicated between ranks (optionally: FSDP)\n",
        "- Each rank holds a subset of sequence tokens\n",
        "- Embeddings, logits, normalizations, MLP all apply independently to token shards\n",
        "- The multi-head attention is the only layer that gets special treatment\n",
        "    - First, apply QKV projections to local tokens, as in data-parallel training;\n",
        "    - Then re-shard so that each rank holds a **subset of heads** across **all tokens**;\n",
        "    - Compute the attention ''core'' (RoPE and F.scaled_dot_product_attention) for its chunk of heads independently;\n",
        "    - Re-shard outputs again so that each rank concatenates **all heads**, but only for its **subset of tokens**;\n",
        "    - Apply the output (\"O\") projection to your local tokens again.\n",
        "- This approach *may* be combined with tensor parallelism, but this is an advanced technique that you don't have to implement.\n",
        "\n",
        "\n",
        "__You have a choice__ between two options on how to implement it: either manually with torch.distributed like in task 2, or using the DTensor route like in task 3. We provide some tips for both tasks.\n",
        "\n",
        "\n",
        "**Option A. with raw `torch.distirbuted`:**\n",
        "- Use [`dist.all_to_all`](https://pytorch.org/docs/stable/distributed.html#torch.distributed.all_to_all) to switch between per-token and per-head sharding without materializing the full tensor on any device;\n",
        "- Wrap the model with [`DistributedDataParallel`](https://pytorch.org/tutorials/intermediate/ddp_tutorial.html) or [`FullyShardedDataParallel`](https://pytorch.org/docs/stable/fsdp.html) so that fine-tuning synchronizes trainable parameters. Note that using FSDP for parameter-efficient fine-tuning can be tricky: we recommend you either wrap **trainable modules** with separate FSDP sub-instances via auto_wrap_policy - or simply use DDP instead of FSDP.\n",
        "\n",
        "**Option B. with `DTensor`:**\n",
        "- We recommend you first skim the official [tutorial](https://pytorch.org/tutorials/intermediate/TP_tutorial.html) on applying Tensor Parallelism (sic.) - or browse the [TorchTitan's version](https://github.com/pytorch/torchtitan/blob/82afc842e303e49d1a137fc7ea48291a57f72d5d/torchtitan/models/llama/parallelize_llama.py) of it.\n",
        "- Note that there is a [`SequenceParallel`](https://pytorch.org/docs/stable/distributed.tensor.parallel.html#torch.distributed.tensor.parallel.SequenceParallel) class in torch.distributed.tensor.parallel` - **however, it does not magick the sequence parallelism for you** - it is only meant for small layers (e.g. normalization). You still need to do the sharding in self-attention!\n",
        "\n",
        "For the sake of formality, here's an example script you need to parallelize:"
      ],
      "metadata": {
        "id": "IWKuXKS3jsLY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import transformers\n",
        "import peft\n",
        "MODEL_NAME = \"unsloth/Llama-3.2-1B\"  # for testing (but not grading!), you may want to use Maykeye/TinyLLama-v0\n",
        "SEQUENCE_LENGTH = 128                # IMPORTANT!!! you need to increase this parameter! Look for the maximum sequence length on one and multiple GPUs\n",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME)\n",
        "model = transformers.LlamaForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.bfloat16).to(device)\n",
        "\n",
        "for param in model.parameters():\n",
        "  param.required_grad = False\n",
        "model.gradient_checkpointing_enable()\n",
        "model.enable_input_require_grads()\n",
        "\n",
        "model = peft.get_peft_model(model, peft.PromptTuningConfig(task_type=peft.TaskType.CAUSAL_LM, num_virtual_tokens=32))\n",
        "assert any(param.requires_grad for param in model.parameters()), \"No trainable parameters - did you enable PEFT?\"\n",
        "\n",
        "!wget -q https://www.gutenberg.org/cache/epub/4300/pg4300.txt -O ulysses.txt  # ... or use any other text of your choosing\n",
        "input_ids = tokenizer(open(\"ulysses.txt\").read(), return_tensors='pt')['input_ids']\n",
        "print(f\"Cropping {input_ids.shape[1]=} to {SEQUENCE_LENGTH} tokens\")\n",
        "input_ids, labels = input_ids[:, :SEQUENCE_LENGTH], input_ids[:, 1:SEQUENCE_LENGTH + 1]\n",
        "\n",
        "trainable_parameters = {p for p in model.parameters() if p.requires_grad}\n",
        "print(f\"Parameters: {sum(map(torch.Tensor.numel, trainable_parameters))} trainable / {sum(map(torch.Tensor.numel, model.parameters()))} total\")\n",
        "opt = torch.optim.Adam(trainable_parameters)\n",
        "for i in range(10):\n",
        "  loss = model(input_ids=input_ids.to(device), labels=labels.to(device)).loss\n",
        "  opt.zero_grad()\n",
        "  loss.backward()\n",
        "  opt.step()\n",
        "  print(f\"{i=}\\t{loss.item()=}\")\n",
        "\n",
        "# pro tip: delete the model or restart session to free RAM for the TP experiments"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 539,
          "referenced_widgets": [
            "bdfc71a449484eaeb3539fabb60a767c",
            "c5a57f756d474bdea5ce6b012edf1508",
            "ce7ec364252f46b996ad0b32fe9d38cd",
            "8d369b7163074ebd822d6702878256a4",
            "ab8fc563343e4e7ea09356642de10756",
            "6a9ff0b43f794891ba1f5c4ecd6dcfdf",
            "3787e2bdb8cd42a8bf1815a090ed9706",
            "b944e375f23d4ea98e58533f11857d33",
            "44d7e3d82cc54756a70a287bbd39ecb9",
            "6f4ac58e5676482e9bcb949adc4f84d1",
            "d2f335cbff4e4ff2baa02a2ad7b82a3e",
            "1aa8ab0872514d57b0592ba43ea0105a",
            "bc7ba0334162415ba5f27f20deb28247",
            "c0e6c92aaebc4f02a4943a572d8220dd",
            "db12af0c8a0d41fa981dbb66e7ed49d4",
            "4b7c453fb3994045a855b7370eee234f",
            "c6e968957e8c4a21953a2d3b9ddd4f8d",
            "2ac894dade0c48b7a30c058eded65b73",
            "94f4a642365b4e69aa15e7ea2bd1a29b",
            "fa7d727b36b742419a928e271d403bba",
            "6ee6f803e8a34061b52b09a1ed9c3a29",
            "dcd4b9306a21422abdd2073913493326",
            "d6509faf9220437e9371f83b5058fc29",
            "d8e556d9c6c04603905217901cb38ac7",
            "820a152271204ba18f6d4808e9fca90e",
            "55e4bf33a30e400db72d47570e2345f8",
            "5ca5324606e84da48635840a32e94f61",
            "6a6840053fad4785a2124aea7c261f59",
            "f1c7bf7153f14344a6e4a7a9bdc35b53",
            "b26d8fd37493485a9c69a20de99262bf",
            "76cbb91b2815478ea29d59b9e4bbc937",
            "7b69fc37e8054da1977eebf3ab652157",
            "0f03319b7e6f4fcc93e2b91ca3293584",
            "8d891195bee042c5b3c88411367196ca",
            "01dd79c9b6404d3c98affd6666ca2d22",
            "50b1790060e94fab95d79317cc0dc333",
            "f9649cfc2d54427ca3a70f87a82f4a5f",
            "57f7eea22c934910b8897d340edf3107",
            "fff708ec9f324f8980f8549639dae9ee",
            "92c3644edc2c48aca666dc94b267dda8",
            "1c80c6a948d84ab2810ea45af8d27af2",
            "74c395453ab64b0899505f152a12c50a",
            "94a2a7b26e984778a2186a491bf06b4f",
            "a7c591494d654cdd8a19141cf96d8036",
            "d77be2d56f3846d4991c92ee953731ec",
            "f940bf4b3d454ad7a3592b0cc51979fd",
            "8339128e688d4221bc318b45910404a6",
            "0309561fa3734a0ea3cd3421c420ff56",
            "2329d7fe319b4383ab87aabe5e3e17fa",
            "df6158bafdd94f4e8b9f147acc0ef8f5",
            "48669f5193124c7ca52d554a998be789",
            "a12be7a58f764dff96d41ba35d3a9c76",
            "ad6a69b6a4614050a497eccebdc920b2",
            "37919cdcd7304db2a60df1cc387ca4ce",
            "d1cbd32f079042398389dd2f250cc8ef",
            "0b353dfbb9784b3883ff31ed7124f450",
            "443a315962ab46ea8b6d8f2ed354038e",
            "d1d68bb4f42647e39806f1fc4dd9c014",
            "ede69839385f48399d51b0c9a0568353",
            "5768b9b7ffec4884869b2334eb946015",
            "9045df041cac425fa7e109d303bb4d2f",
            "0dc89a0bd4e34ce1b60597cb1acd7fdf",
            "5a5c1dcd5a214c93bb917c8eb9029af2",
            "c5f9ad4f3dd94e8b8d05530f11fed714",
            "fd64196b94f64bf59d4bbd1307aaf05c",
            "d643e91490274b1988b242e4d784ff82"
          ]
        },
        "id": "9ofi1_Kgusd8",
        "outputId": "d456bd60-be30-4f81-917d-5ef3d9e04538"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.11/dist-packages/huggingface_hub/utils/_auth.py:94: UserWarning: \n",
            "The secret `HF_TOKEN` does not exist in your Colab secrets.\n",
            "To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session.\n",
            "You will be able to reuse this secret in all of your notebooks.\n",
            "Please note that authentication is recommended but still optional to access public models or datasets.\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "tokenizer_config.json:   0%|          | 0.00/50.6k [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "bdfc71a449484eaeb3539fabb60a767c"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "tokenizer.json:   0%|          | 0.00/17.2M [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "1aa8ab0872514d57b0592ba43ea0105a"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "special_tokens_map.json:   0%|          | 0.00/459 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "d6509faf9220437e9371f83b5058fc29"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "config.json:   0%|          | 0.00/935 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "8d891195bee042c5b3c88411367196ca"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "model.safetensors:   0%|          | 0.00/2.47G [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "d77be2d56f3846d4991c92ee953731ec"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "generation_config.json:   0%|          | 0.00/230 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "0b353dfbb9784b3883ff31ed7124f450"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Token indices sequence length is longer than the specified maximum sequence length for this model (397368 > 131072). Running this sequence through the model will result in indexing errors\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cropping input_ids.shape[1]=397368 to 128 tokens\n",
            "Parameters: 65536 trainable / 1235879936 total\n",
            "i=0\tloss.item()=8.582364082336426\n",
            "i=1\tloss.item()=8.413138389587402\n",
            "i=2\tloss.item()=8.251019477844238\n",
            "i=3\tloss.item()=8.113922119140625\n",
            "i=4\tloss.item()=8.001195907592773\n",
            "i=5\tloss.item()=7.875072002410889\n",
            "i=6\tloss.item()=7.784201145172119\n",
            "i=7\tloss.item()=7.691455364227295\n",
            "i=8\tloss.item()=7.621612548828125\n",
            "i=9\tloss.item()=7.547983169555664\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "__Task 4 (1 point):__ before you do training, let's first parallelize a single forward pass. Implement sharding with the same interface you used in tasks 2 (or 3 if you use DTensor), but this time, parallelize across the sequence dimension. Note: if you are running out of (V)RAM, load the 1B model in half precision and disable gradients for all weights except the first (few) layers.\n",
        "\n"
      ],
      "metadata": {
        "id": "kelsoIRmtyPL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "<<A whole lot of your code here>>"
      ],
      "metadata": {
        "id": "2xwwjXfPdnqt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "<<... and a dedicated cell to show off that it works>>"
      ],
      "metadata": {
        "id": "a7Vw7iH-kLMc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "__Task 5 (1 point):__ Now use the script above to parallelize the entire training run. You are free to use other fine-tuning methods (e.g. LoRA or even full fine-tuning), as long as you can demonstrate that the loss goes down.\n",
        "\n",
        "**Make sure you increase SEQUENCE_LENGTH as much as possible!** Even on a single GPU, you should be able to go into thousands, if not tens of thousands of tokens - and report the maximum sequence length with one and with multiple GPUs respectively.\n",
        "\n",
        "If you don't have access to multiple GPUs, you may optionally submit a version that does training on a single GPU, but computes attention heads sequentially with gradient checkpointing - but if you do, please announce that you are using this option in bold, capital letters, so that the grader will notice it."
      ],
      "metadata": {
        "id": "oS4JWk3J7oP9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "<<A whole lot of your code here>>"
      ],
      "metadata": {
        "id": "azcZBvaa8ymE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "<<... and a dedicated cell to show off that it works>>"
      ],
      "metadata": {
        "id": "xBWi66FkC1dz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "```\n",
        "\n",
        "\n",
        "### Optional: bonus tasks\n",
        "\n",
        "There are many routes to further improve the training/inference code. You may (but you don't have to) implement any combination of them for bonus points.\n",
        "\n",
        "However, please not that the total points for this week's entire assignment (part 1 & 2) are **capped at 14**.\n",
        "\n",
        "__Bonus task: parallel key-value caching (1 point).__ In tasks 2 and 3, you implement tensor parallelism for attention forward pass and perform inference with re-computation. However, real world inference engines use [KV caching](https://huggingface.co/docs/transformers/main/en/kv_cache) - keeping key and value caches from past tokens and only processing the new token each time.\n",
        "\n",
        "For this task, you will have to implement this type of parallelism for either torch.distributed or DTensor implementation of attention $-$ simply cache the heads already assigned to each rank. To get the grade, you will need to demonstrate that the model generates a sensible text with any cache (via past_key_values=).\n",
        "\n",
        "__Bonus task: pipeline parallelism (1-2 points):__ In tasks 1-3, you've implemented symmetric model parallelism, aka Tensor Parallelism. However, there is another way to partition model parameters $-$ assign entire layers to each rank and run them in a pipeline. This can be faster, especially if you are running\n",
        "\n",
        "For 1 point, check out [torch.distributed.pipelinging](https://pytorch.org/docs/stable/distributed.pipelining.html), [DeepSpeed pipelining](https://deepspeed.readthedocs.io/en/latest/pipeline.html) or [torchgpipe](https://github.com/kakaobrain/torchgpipe) and demonstrate that you can run or fine-tune a model that would not fit into a single GPU (you will need multuple devices for this!).\n",
        "\n",
        "For 2 points, compare different pipelining schedules in terms of training throughput: use GPipe as a baseline and try ScheduleInterleaved1F1B (or a more advanced pipeline of your choosing).\n",
        "\n",
        "__Bonus task: better sequence parallelism (2 points).__ In tasks 4 and 5, you implemented basic sequence parallelism. However, there are multiple ways you can improve that technique for further memory savings or better device utilization.\n",
        "\n",
        "For 1 point, implement combined tensor + sequence parallelism and compare results with naive sequence parallelism.\n",
        "\n",
        "For 2 points, implement [Ring Attention](https://arxiv.org/abs/2310.01889) *or* integrate computation-communication overlap from [FLUX](https://arxiv.org/abs/2406.06858) and measure the speed and memory trade-offs."
      ],
      "metadata": {
        "id": "fZFeok82812C"
      }
    }
  ]
}