Upload 17 files

README.md

@@ -0,0 +1,603 @@
+---
+tags:
+- summarization
+- summary
+- booksum
+- long-document
+- long-form
+license:
+- apache-2.0
+- bsd-3-clause
+datasets:
+- kmfoda/booksum
+metrics:
+- rouge
+widget:
+- text: large earthquakes along a given fault segment do not occur at random intervals
+    because it takes time to accumulate the strain energy for the rupture. The rates
+    at which tectonic plates move and accumulate strain at their boundaries are approximately
+    uniform. Therefore, in first approximation, one may expect that large ruptures
+    of the same fault segment will occur at approximately constant time intervals.
+    If subsequent main shocks have different amounts of slip across the fault, then
+    the recurrence time may vary, and the basic idea of periodic mainshocks must be
+    modified. For great plate boundary ruptures the length and slip often vary by
+    a factor of 2. Along the southern segment of the San Andreas fault the recurrence
+    interval is 145 years with variations of several decades. The smaller the standard
+    deviation of the average recurrence interval, the more specific could be the long
+    term prediction of a future mainshock.
+  example_title: earthquakes
+- text: " A typical feed-forward neural field algorithm. Spatiotemporal coordinates\
+    \ are fed into a neural network that predicts values in the reconstructed domain.\
+    \ Then, this domain is mapped to the sensor domain where sensor measurements are\
+    \ available as supervision. Class and Section Problems Addressed Generalization\
+    \ (Section 2) Inverse problems, ill-posed problems, editability; symmetries. Hybrid\
+    \ Representations (Section 3) Computation & memory efficiency, representation\
+    \ capacity, editability: Forward Maps (Section 4) Inverse problems Network Architecture\
+    \ (Section 5) Spectral bias, integration & derivatives. Manipulating Neural Fields\
+    \ (Section 6) Edit ability, constraints, regularization. Table 2: The five classes\
+    \ of techniques in the neural field toolbox each addresses problems that arise\
+    \ in learning, inference, and control. (Section 3). We can supervise reconstruction\
+    \ via differentiable forward maps that transform Or project our domain (e.g, 3D\
+    \ reconstruction via 2D images; Section 4) With appropriate network architecture\
+    \ choices, we can overcome neural network spectral biases (blurriness) and efficiently\
+    \ compute derivatives and integrals (Section 5). Finally, we can manipulate neural\
+    \ fields to add constraints and regularizations, and to achieve editable representations\
+    \ (Section 6). Collectively, these classes constitute a 'toolbox' of techniques\
+    \ to help solve problems with neural fields There are three components in a conditional\
+    \ neural field: (1) An encoder or inference function \u20AC that outputs the conditioning\
+    \ latent variable 2 given an observation 0 E(0) =2. 2 is typically a low-dimensional\
+    \ vector, and is often referred to aS a latent code Or feature code_ (2) A mapping\
+    \ function 4 between Z and neural field parameters O: Y(z) = O; (3) The neural\
+    \ field itself $. The encoder \u20AC finds the most probable z given the observations\
+    \ O: argmaxz P(2/0). The decoder maximizes the inverse conditional probability\
+    \ to find the most probable 0 given Z: arg- max P(Olz). We discuss different encoding\
+    \ schemes with different optimality guarantees (Section 2.1.1), both global and\
+    \ local conditioning (Section 2.1.2), and different mapping functions Y (Section\
+    \ 2.1.3) 2. Generalization Suppose we wish to estimate a plausible 3D surface\
+    \ shape given a partial or noisy point cloud. We need a suitable prior over the\
+    \ sur- face in its reconstruction domain to generalize to the partial observations.\
+    \ A neural network expresses a prior via the function space of its architecture\
+    \ and parameters 0, and generalization is influenced by the inductive bias of\
+    \ this function space (Section 5)."
+  example_title: scientific paper
+- text: 'Is a else or outside the cob and tree written being of early client rope
+    and you have is for good reasons. On to the ocean in Orange for time. By''s the
+    aggregate we can bed it yet. Why this please pick up on a sort is do and also
+    M Getoi''s nerocos and do rain become you to let so is his brother is made in
+    use and Mjulia''s''s the lay major is aging Masastup coin present sea only of
+    Oosii rooms set to you We do er do we easy this private oliiishs lonthen might
+    be okay. Good afternoon everybody. Welcome to this lecture of Computational Statistics.
+    As you can see, I''m not socially my name is Michael Zelinger. I''m one of the
+    task for this class and you might have already seen me in the first lecture where
+    I made a quick appearance. I''m also going to give the tortillas in the last third
+    of this course. So to give you a little bit about me, I''m a old student here
+    with better Bulman and my research centres on casual inference applied to biomedical
+    disasters, so that could be genomics or that could be hospital data. If any of
+    you is interested in writing a bachelor thesis, a semester paper may be mastathesis
+    about this topic feel for reach out to me. you have my name on models and my email
+    address you can find in the directory I''d Be very happy to talk about it. you
+    do not need to be sure about it, we can just have a chat. So with that said, let''s
+    get on with the lecture. There''s an exciting topic today I''m going to start
+    by sharing some slides with you and later on during the lecture we''ll move to
+    the paper. So bear with me for a few seconds. Well, the projector is starting
+    up. Okay, so let''s get started. Today''s topic is a very important one. It''s
+    about a technique which really forms one of the fundamentals of data science,
+    machine learning, and any sort of modern statistics. It''s called cross validation.
+    I know you really want to understand this topic I Want you to understand this
+    and frankly, nobody''s gonna leave Professor Mineshousen''s class without understanding
+    cross validation. So to set the stage for this, I Want to introduce you to the
+    validation problem in computational statistics. So the problem is the following:
+    You trained a model on available data. You fitted your model, but you know the
+    training data you got could always have been different and some data from the
+    environment. Maybe it''s a random process. You do not really know what it is,
+    but you know that somebody else who gets a different batch of data from the same
+    environment they would get slightly different training data and you do not care
+    that your method performs as well. On this training data. you want to to perform
+    well on other data that you have not seen other data from the same environment.
+    So in other words, the validation problem is you want to quantify the performance
+    of your model on data that you have not seen. So how is this even possible? How
+    could you possibly measure the performance on data that you do not know The solution
+    to? This is the following realization is that given that you have a bunch of data,
+    you were in charge. You get to control how much that your model sees. It works
+    in the following way: You can hide data firms model. Let''s say you have a training
+    data set which is a bunch of doubtless so X eyes are the features those are typically
+    hide and national vector. It''s got more than one dimension for sure. And the
+    why why eyes. Those are the labels for supervised learning. As you''ve seen before,
+    it''s the same set up as we have in regression. And so you have this training
+    data and now you choose that you only use some of those data to fit your model.
+    You''re not going to use everything, you only use some of it the other part you
+    hide from your model. And then you can use this hidden data to do validation from
+    the point of you of your model. This hidden data is complete by unseen. In other
+    words, we solve our problem of validation.'
+  example_title: transcribed audio - lecture
+- text: "Transformer-based models have shown to be very useful for many NLP tasks.\
+    \ However, a major limitation of transformers-based models is its O(n^2)O(n 2)\
+    \ time & memory complexity (where nn is sequence length). Hence, it's computationally\
+    \ very expensive to apply transformer-based models on long sequences n > 512n>512.\
+    \ Several recent papers, e.g. Longformer, Performer, Reformer, Clustered attention\
+    \ try to remedy this problem by approximating the full attention matrix. You can\
+    \ checkout \U0001F917's recent blog post in case you are unfamiliar with these\
+    \ models.\nBigBird (introduced in paper) is one of such recent models to address\
+    \ this issue. BigBird relies on block sparse attention instead of normal attention\
+    \ (i.e. BERT's attention) and can handle sequences up to a length of 4096 at a\
+    \ much lower computational cost compared to BERT. It has achieved SOTA on various\
+    \ tasks involving very long sequences such as long documents summarization, question-answering\
+    \ with long contexts.\nBigBird RoBERTa-like model is now available in \U0001F917\
+    Transformers. The goal of this post is to give the reader an in-depth understanding\
+    \ of big bird implementation & ease one's life in using BigBird with \U0001F917\
+    Transformers. But, before going into more depth, it is important to remember that\
+    \ the BigBird's attention is an approximation of BERT's full attention and therefore\
+    \ does not strive to be better than BERT's full attention, but rather to be more\
+    \ efficient. It simply allows to apply transformer-based models to much longer\
+    \ sequences since BERT's quadratic memory requirement quickly becomes unbearable.\
+    \ Simply put, if we would have \u221E compute & \u221E time, BERT's attention\
+    \ would be preferred over block sparse attention (which we are going to discuss\
+    \ in this post).\nIf you wonder why we need more compute when working with longer\
+    \ sequences, this blog post is just right for you!\nSome of the main questions\
+    \ one might have when working with standard BERT-like attention include:\nDo all\
+    \ tokens really have to attend to all other tokens? Why not compute attention\
+    \ only over important tokens? How to decide what tokens are important? How to\
+    \ attend to just a few tokens in a very efficient way? In this blog post, we will\
+    \ try to answer those questions.\nWhat tokens should be attended to? We will give\
+    \ a practical example of how attention works by considering the sentence 'BigBird\
+    \ is now available in HuggingFace for extractive question answering'. In BERT-like\
+    \ attention, every word would simply attend to all other tokens.\nLet's think\
+    \ about a sensible choice of key tokens that a queried token actually only should\
+    \ attend to by writing some pseudo-code. Will will assume that the token available\
+    \ is queried and build a sensible list of key tokens to attend to.\n>>> # let's\
+    \ consider following sentence as an example >>> example = ['BigBird', 'is', 'now',\
+    \ 'available', 'in', 'HuggingFace', 'for', 'extractive', 'question', 'answering']\n\
+    >>> # further let's assume, we're trying to understand the representation of 'available'\
+    \ i.e. >>> query_token = 'available' >>> # We will initialize an empty `set` and\
+    \ fill up the tokens of our interest as we proceed in this section. >>> key_tokens\
+    \ = [] # => currently 'available' token doesn't have anything to attend Nearby\
+    \ tokens should be important because, in a sentence (sequence of words), the current\
+    \ word is highly dependent on neighboring past & future tokens. This intuition\
+    \ is the idea behind the concept of sliding attention."
+  example_title: bigbird blog intro
+- text: "To be fair, you have to have a very high IQ to understand Rick and Morty.\
+    \ The humour is extremely subtle, and without a solid grasp of theoretical physics\
+    \ most of the jokes will go over a typical viewer's head. There's also Rick's\
+    \ nihilistic outlook, which is deftly woven into his characterisation- his personal\
+    \ philosophy draws heavily from Narodnaya Volya literature, for instance. The\
+    \ fans understand this stuff; they have the intellectual capacity to truly appreciate\
+    \ the depths of these jokes, to realise that they're not just funny- they say\
+    \ something deep about LIFE. As a consequence people who dislike Rick & Morty\
+    \ truly ARE idiots- of course they wouldn't appreciate, for instance, the humour\
+    \ in Rick's existential catchphrase 'Wubba Lubba Dub Dub,' which itself is a cryptic\
+    \ reference to Turgenev's Russian epic Fathers and Sons. I'm smirking right now\
+    \ just imagining one of those addlepated simpletons scratching their heads in\
+    \ confusion as Dan Harmon's genius wit unfolds itself on their television screens.\
+    \ What fools.. how I pity them. \U0001F602\nAnd yes, by the way, i DO have a Rick\
+    \ & Morty tattoo. And no, you cannot see it. It's for the ladies' eyes only- and\
+    \ even then they have to demonstrate that they're within 5 IQ points of my own\
+    \ (preferably lower) beforehand. Nothin personnel kid \U0001F60E"
+  example_title: Richard & Mortimer
+- text: "The tower is 324 metres (1,063 ft) tall, about the same height as an 81-storey building, and the tallest structure in Paris. Its base is square, measuring 125 metres (410 ft) on each side. During its construction, the Eiffel Tower surpassed the Washington Monument to become the tallest man-made structure in the world, a title it held for 41 years until the Chrysler Building in New York City was finished in 1930. It was the first structure to reach a height of 300 metres. Due to the addition of a broadcasting aerial at the top of the tower in 1957, it is now taller than the Chrysler Building by 5.2 metres (17 ft). Excluding transmitters, the Eiffel Tower is the second tallest free-standing structure in France after the Millau Viaduct."
+  example_title: eiffel
+parameters:
+  max_length: 64
+  min_length: 8
+  no_repeat_ngram_size: 3
+  early_stopping: true
+  repetition_penalty: 3.5
+  length_penalty: 0.3
+  encoder_no_repeat_ngram_size: 3
+  num_beams: 4
+model-index:
+- name: pszemraj/long-t5-tglobal-base-16384-book-summary
+  results:
+  - task:
+      type: summarization
+      name: Summarization
+    dataset:
+      name: kmfoda/booksum
+      type: kmfoda/booksum
+      config: kmfoda--booksum
+      split: test
+    metrics:
+    - name: ROUGE-1
+      type: rouge
+      value: 36.4085
+      verified: true
+    - name: ROUGE-2
+      type: rouge
+      value: 6.0646
+      verified: true
+    - name: ROUGE-L
+      type: rouge
+      value: 16.7209
+      verified: true
+    - name: ROUGE-LSUM
+      type: rouge
+      value: 33.3405
+      verified: true
+    - name: loss
+      type: loss
+      value: .nan
+      verified: true
+    - name: gen_len
+      type: gen_len
+      value: 252.8099
+      verified: true
+  - task:
+      type: summarization
+      name: Summarization
+    dataset:
+      name: samsum
+      type: samsum
+      config: samsum
+      split: test
+    metrics:
+    - name: ROUGE-1
+      type: rouge
+      value: 30.9047
+      verified: true
+    - name: ROUGE-2
+      type: rouge
+      value: 7.4715
+      verified: true
+    - name: ROUGE-L
+      type: rouge
+      value: 22.3962
+      verified: true
+    - name: ROUGE-LSUM
+      type: rouge
+      value: 26.9094
+      verified: true
+    - name: loss
+      type: loss
+      value: .nan
+      verified: true
+    - name: gen_len
+      type: gen_len
+      value: 46.7973
+      verified: true
+  - task:
+      type: summarization
+      name: Summarization
+    dataset:
+      name: cnn_dailymail
+      type: cnn_dailymail
+      config: 3.0.0
+      split: test
+    metrics:
+    - name: ROUGE-1
+      type: rouge
+      value: 30.5942
+      verified: true
+    - name: ROUGE-2
+      type: rouge
+      value: 7.252
+      verified: true
+    - name: ROUGE-L
+      type: rouge
+      value: 17.7156
+      verified: true
+    - name: ROUGE-LSUM
+      type: rouge
+      value: 27.2881
+      verified: true
+    - name: loss
+      type: loss
+      value: .nan
+      verified: true
+    - name: gen_len
+      type: gen_len
+      value: 125.2507
+      verified: true
+  - task:
+      type: summarization
+      name: Summarization
+    dataset:
+      name: xsum
+      type: xsum
+      config: default
+      split: test
+    metrics:
+    - name: ROUGE-1
+      type: rouge
+      value: 20.3648
+      verified: true
+    - name: ROUGE-2
+      type: rouge
+      value: 3.4126
+      verified: true
+    - name: ROUGE-L
+      type: rouge
+      value: 13.6168
+      verified: true
+    - name: ROUGE-LSUM
+      type: rouge
+      value: 15.8313
+      verified: true
+    - name: loss
+      type: loss
+      value: .nan
+      verified: true
+    - name: gen_len
+      type: gen_len
+      value: 82.2177
+      verified: true
+  - task:
+      type: summarization
+      name: Summarization
+    dataset:
+      name: billsum
+      type: billsum
+      config: default
+      split: test
+    metrics:
+    - name: ROUGE-1
+      type: rouge
+      value: 39.6378
+      verified: true
+    - name: ROUGE-2
+      type: rouge
+      value: 13.0017
+      verified: true
+    - name: ROUGE-L
+      type: rouge
+      value: 23.0255
+      verified: true
+    - name: ROUGE-LSUM
+      type: rouge
+      value: 32.9943
+      verified: true
+    - name: loss
+      type: loss
+      value: 1.9428048133850098
+      verified: true
+    - name: gen_len
+      type: gen_len
+      value: 162.3588
+      verified: true
+  - task:
+      type: summarization
+      name: Summarization
+    dataset:
+      name: big_patent
+      type: big_patent
+      config: y
+      split: test
+    metrics:
+    - name: ROUGE-1
+      type: rouge
+      value: 34.7641
+      verified: true
+    - name: ROUGE-2
+      type: rouge
+      value: 7.8744
+      verified: true
+    - name: ROUGE-L
+      type: rouge
+      value: 19.9826
+      verified: true
+    - name: ROUGE-LSUM
+      type: rouge
+      value: 29.208
+      verified: true
+    - name: loss
+      type: loss
+      value: 2.8316469192504883
+      verified: true
+    - name: gen_len
+      type: gen_len
+      value: 132.7475
+      verified: true
+  - task:
+      type: summarization
+      name: Summarization
+    dataset:
+      name: launch/gov_report
+      type: launch/gov_report
+      config: plain_text
+      split: validation
+    metrics:
+    - name: ROUGE-1
+      type: rouge
+      value: 37.9246
+      verified: true
+    - name: ROUGE-2
+      type: rouge
+      value: 8.5837
+      verified: true
+    - name: ROUGE-L
+      type: rouge
+      value: 18.0274
+      verified: true
+    - name: ROUGE-LSUM
+      type: rouge
+      value: 34.0816
+      verified: true
+    - name: loss
+      type: loss
+      value: 2.56695818901062
+      verified: true
+    - name: gen_len
+      type: gen_len
+      value: 220.3747
+      verified: true
+  - task:
+      type: summarization
+      name: Summarization
+    dataset:
+      name: launch/gov_report
+      type: launch/gov_report
+      config: plain_text
+      split: test
+    metrics:
+    - name: ROUGE-1
+      type: rouge
+      value: 37.4438
+      verified: true
+    - name: ROUGE-2
+      type: rouge
+      value: 8.2907
+      verified: true
+    - name: ROUGE-L
+      type: rouge
+      value: 17.6893
+      verified: true
+    - name: ROUGE-LSUM
+      type: rouge
+      value: 33.7141
+      verified: true
+    - name: loss
+      type: loss
+      value: 2.5776000022888184
+      verified: true
+    - name: gen_len
+      type: gen_len
+      value: 214.9692
+      verified: true
+---
+
+# long-t5-tglobal-base-16384 + BookSum
+
+ <a href="https://colab.research.google.com/gist/pszemraj/d9a0495861776168fd5cdcd7731bc4ee/example-long-t5-tglobal-base-16384-book-summary.ipynb">
+  <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
+</a>
+
+Summarize long text and get a SparkNotes-esque summary of arbitrary topics!
+
+- generalizes reasonably well to academic & narrative text.
+- A simple example/use case on ASR is [here](https://longt5-booksum-example.netlify.app/).
+- Example notebook in Colab (_click on the icon above_).
+
+## Cheeky Proof-of-Concept
+
+A summary of the [infamous navy seals copypasta](https://knowyourmeme.com/memes/navy-seal-copypasta):
+
+> The narrator tells us that he's graduated from the Navy seals and has been involved in many secret raids. He's also one of the best snipers in the entire U.S. military. He promises to "wipe you out with precision" when they meet again.
+
+* * *
+
+**Contents**
+
+<!-- TOC -->
+
+- [Model description](#model-description)
+- [How-To in Python](#how-to-in-python)
+- [Intended uses & limitations](#intended-uses--limitations)
+- [Training and evaluation data](#training-and-evaluation-data)
+- [FAQ](#faq)
+    - [Inference over long documents in batches](#how-to-run-inference-over-a-very-long-30k-tokens-document-in-batches)
+    - [How to fine-tune further](#how-to-fine-tune-further)
+- [Training procedure](#training-procedure)
+    - [Updates](#updates)
+    - [Training hyperparameters](#training-hyperparameters)
+    - [Framework versions](#framework-versions)
+- [Citation info](#citation-info)
+
+<!-- /TOC -->
+
+* * *
+
+## Model description
+
+A fine-tuned version of [google/long-t5-tglobal-base](https://huggingface.co/google/long-t5-tglobal-base) on the `kmfoda/booksum` dataset:
+
+- 30+ epochs of fine-tuning from the base model on V100/A100 GPUs
+- Training used 16384 token input / 1024 max output
+
+Read the paper by Guo et al. here: [LongT5: Efficient Text-To-Text Transformer for Long Sequences](https://arxiv.org/pdf/2112.07916.pdf)
+
+## How-To in Python
+
+Install/update transformers `pip install -U transformers`
+
+Summarize text with pipeline:
+
+```python
+import torch
+from transformers import pipeline
+
+summarizer = pipeline(
+    "summarization",
+    "pszemraj/long-t5-tglobal-base-16384-book-summary",
+    device=0 if torch.cuda.is_available() else -1,
+)
+long_text = "Here is a lot of text I don't want to read. Replace me"
+
+result = summarizer(long_text)
+print(result[0]["summary_text"])
+```
+
+Pass [other parameters related to beam search textgen](https://huggingface.co/blog/how-to-generate) when calling `summarizer` to get even higher quality results.
+
+## Intended uses & limitations
+
+- The current checkpoint is fairly well converged but will be updated if further improvements can be made.
+    - Compare performance to [LED-base](https://huggingface.co/pszemraj/led-base-book-summary) trained on the same dataset (API gen parameters are the same).
+- while this model seems to improve upon factual consistency, **do not take summaries to be foolproof and check things that seem odd**.
+
+## Training and evaluation data
+
+`kmfoda/booksum` dataset on HuggingFace - read [the original paper here](https://arxiv.org/abs/2105.08209). Summaries longer than 1024 LongT5 tokens were filtered out to prevent the model from learning to generate "partial" summaries.
+
+
+* * *
+
+## FAQ
+
+### How to run inference over a very long (30k+ tokens) document in batches?
+
+See `summarize.py` in [the code for my hf space Document Summarization](https://huggingface.co/spaces/pszemraj/document-summarization/blob/main/summarize.py) :)
+
+You can also use the same code to split a document into batches of 4096, etc., and run over those with the model. This is useful in situations where CUDA memory is limited.
+
+### How to fine-tune further?
+
+See [train with a script](https://huggingface.co/docs/transformers/run_scripts) and [the summarization scripts](https://github.com/huggingface/transformers/tree/main/examples/pytorch/summarization).
+
+This model was originally tuned on Google Colab with a heavily modified variant of the [longformer training notebook](https://github.com/patrickvonplaten/notebooks/blob/master/Fine_tune_Longformer_Encoder_Decoder_(LED)_for_Summarization_on_pubmed.ipynb), key enabler being deepspeed. You can try this as an alternate route to fine-tuning the model without using the command line.
+
+* * *
+
+## Training procedure
+
+### Updates:
+
+- July 22, 2022: updated to a fairly converged checkpoint
+- July 3, 2022: Added a new version with several epochs of additional general training that is more performant.
+
+### Training hyperparameters
+
+_NOTE: early checkpoints of this model were trained on a "smaller" subsection of the dataset as it was filtered for summaries of **1024 characters**. This was subsequently caught and adjusted to **1024 tokens** and then trained further for 10+ epochs._
+
+The following hyperparameters were used during the **most recent** training round\*:
+
+- learning_rate: 0.0005
+- train_batch_size: 1
+- eval_batch_size: 1
+- seed: 42
+- distributed_type: multi-GPU
+- gradient_accumulation_steps: 128
+- total_train_batch_size: 128
+- optimizer: Adam with betas=(0.9,0.999) and epsilon=1e-08
+- lr_scheduler_type: cosine
+- lr_scheduler_warmup_ratio: 0.01
+- num_epochs: 2
+
+\* Prior training sessions used roughly similar parameters; multiple sessions were required as this takes eons to train
+
+### Framework versions
+
+- Transformers 4.20.1
+- Pytorch 1.10.0+cu113
+- Datasets 2.3.2
+- Tokenizers 0.12.1
+
+## Citation info
+
+If you find `pszemraj/long-t5-tglobal-base-16384-book-summary` useful in your work, please consider citing this model :)
+
+    @misc {peter_szemraj_2022,
+    	author       = { {Peter Szemraj} },
+    	title        = { long-t5-tglobal-base-16384-book-summary (Revision 4b12bce) },
+    	year         = 2022,
+    	url          = { https://huggingface.co/pszemraj/long-t5-tglobal-base-16384-book-summary },
+    	doi          = { 10.57967/hf/0100 },
+    	publisher    = { Hugging Face }
+    }

config.json

@@ -0,0 +1,43 @@
+{
+  "_name_or_path": "pszemraj/long-t5-tglobal-base-16384-booksum-V10",
+  "architectures": [
+    "LongT5ForConditionalGeneration"
+  ],
+  "d_ff": 2048,
+  "d_kv": 64,
+  "d_model": 768,
+  "decoder_start_token_id": 0,
+  "dense_act_fn": "gelu_new",
+  "dropout_rate": 0.1,
+  "early_stopping": true,
+  "encoder_attention_type": "transient-global",
+  "encoder_no_repeat_ngram_size": 4,
+  "eos_token_id": 1,
+  "feed_forward_proj": "gated-gelu",
+  "global_block_size": 16,
+  "initializer_factor": 1.0,
+  "is_encoder_decoder": true,
+  "is_gated_act": true,
+  "layer_norm_epsilon": 1e-06,
+  "length_penalty": 0.8,
+  "local_radius": 127,
+  "max_length": 512,
+  "min_length": 8,
+  "model_type": "longt5",
+  "n_positions": 4096,
+  "no_repeat_ngram_size": 3,
+  "num_beams": 2,
+  "num_decoder_layers": 12,
+  "num_heads": 12,
+  "num_layers": 12,
+  "output_past": true,
+  "pad_token_id": 0,
+  "relative_attention_max_distance": 128,
+  "relative_attention_num_buckets": 32,
+  "repetition_penalty": 3.5,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.20.1",
+  "use_cache": false,
+  "vocab_size": 32128
+}

index.gitattributes

@@ -0,0 +1,27 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

index.gitignore

@@ -0,0 +1 @@
+checkpoint-*/

latest

@@ -0,0 +1 @@
+global_step330

long-t5-tglobal-base-16384-booksum-V6-partial-ft3-booksum_training_metadata.json

@@ -0,0 +1 @@
+{"output_dir": "/content/drive/MyDrive/Programming/hf-trainer/long-t5-tglobal-base-16384-booksum-V6-partial-ft3-booksum", "overwrite_output_dir": true, "do_train": false, "do_eval": false, "do_predict": false, "evaluation_strategy": "no", "prediction_loss_only": false, "per_device_train_batch_size": 2, "per_device_eval_batch_size": 1, "per_gpu_train_batch_size": "None", "per_gpu_eval_batch_size": "None", "gradient_accumulation_steps": 64, "eval_accumulation_steps": "None", "eval_delay": 0, "learning_rate": 0.0004, "weight_decay": 0.05, "adam_beta1": 0.9, "adam_beta2": 0.999, "adam_epsilon": 1e-08, "max_grad_norm": 0.2, "num_train_epochs": 2, "max_steps": -1, "lr_scheduler_type": "cosine_with_restarts", "warmup_ratio": 0.02, "warmup_steps": 0, "log_level": -1, "log_level_replica": -1, "log_on_each_node": true, "logging_dir": "/content/drive/MyDrive/Programming/hf-trainer/long-t5-tglobal-base-16384-booksum-V6-partial-ft3-booksum/logs", "logging_strategy": "steps", "logging_first_step": false, "logging_steps": 3, "logging_nan_inf_filter": true, "save_strategy": "epoch", "save_steps": 500, "save_total_limit": 1, "save_on_each_node": false, "no_cuda": false, "seed": 42, "data_seed": "None", "jit_mode_eval": false, "use_ipex": false, "bf16": false, "fp16": false, "fp16_opt_level": "O1", "half_precision_backend": "auto", "bf16_full_eval": false, "fp16_full_eval": false, "tf32": "None", "local_rank": 0, "xpu_backend": "None", "tpu_num_cores": "None", "tpu_metrics_debug": false, "debug": "[]", "dataloader_drop_last": false, "eval_steps": "None", "dataloader_num_workers": 0, "past_index": -1, "run_name": "/content/drive/MyDrive/Programming/hf-trainer/long-t5-tglobal-base-16384-booksum-V6-partial-ft3-booksum", "disable_tqdm": false, "remove_unused_columns": true, "label_names": "None", "load_best_model_at_end": false, "metric_for_best_model": "None", "greater_is_better": "None", "ignore_data_skip": false, "sharded_ddp": "[]", "fsdp": "[]", "fsdp_min_num_params": 0, "deepspeed": "/content/ds_config_zero2_bf16.json", "label_smoothing_factor": 0.0, "optim": "adamw_hf", "adafactor": false, "group_by_length": false, "length_column_name": "length", "report_to": "['tensorboard']", "ddp_find_unused_parameters": "None", "ddp_bucket_cap_mb": "None", "dataloader_pin_memory": true, "skip_memory_metrics": true, "use_legacy_prediction_loop": false, "push_to_hub": true, "resume_from_checkpoint": "None", "hub_model_id": "long-t5-tglobal-base-16384-booksum-V6-partial-ft3-booksum", "hub_strategy": "end", "hub_token": "<HUB_TOKEN>", "hub_private_repo": true, "gradient_checkpointing": true, "include_inputs_for_metrics": false, "fp16_backend": "auto", "push_to_hub_model_id": "None", "push_to_hub_organization": "None", "push_to_hub_token": "<PUSH_TO_HUB_TOKEN>", "_n_gpu": 1, "mp_parameters": "", "auto_find_batch_size": false, "full_determinism": false, "torchdynamo": "None", "ray_scope": "last", "sortish_sampler": false, "predict_with_generate": false, "generation_max_length": "None", "generation_num_beams": "None", "train_batch_size": 2, "eval_batch_size": 1, "configs_src": "long-t5-tglobal-base-16384-booksum-V6-partial-ft3-booksum"}

long-t5-tglobal-base-16384-booksum-V7.9-ft1-booksum_training_metadata.json

@@ -0,0 +1 @@
+{"output_dir": "/content/drive/MyDrive/Programming/hf-trainer/long-t5-tglobal-base-16384-booksum-V7.9-ft1-booksum", "overwrite_output_dir": true, "do_train": false, "do_eval": false, "do_predict": false, "evaluation_strategy": "no", "prediction_loss_only": false, "per_device_train_batch_size": 1, "per_device_eval_batch_size": 1, "per_gpu_train_batch_size": "None", "per_gpu_eval_batch_size": "None", "gradient_accumulation_steps": 64, "eval_accumulation_steps": "None", "eval_delay": 0, "learning_rate": 0.001, "weight_decay": 0.05, "adam_beta1": 0.9, "adam_beta2": 0.999, "adam_epsilon": 1e-08, "max_grad_norm": 0.5, "num_train_epochs": 2, "max_steps": -1, "lr_scheduler_type": "cosine", "warmup_ratio": 0.01, "warmup_steps": 0, "log_level": -1, "log_level_replica": -1, "log_on_each_node": true, "logging_dir": "/content/drive/MyDrive/Programming/hf-trainer/long-t5-tglobal-base-16384-booksum-V7.9-ft1-booksum/logs", "logging_strategy": "steps", "logging_first_step": false, "logging_steps": 2, "logging_nan_inf_filter": true, "save_strategy": "epoch", "save_steps": 500, "save_total_limit": 1, "save_on_each_node": false, "no_cuda": false, "seed": 42, "data_seed": "None", "jit_mode_eval": false, "use_ipex": false, "bf16": false, "fp16": true, "fp16_opt_level": "O1", "half_precision_backend": "cuda_amp", "bf16_full_eval": false, "fp16_full_eval": false, "tf32": "None", "local_rank": 0, "xpu_backend": "None", "tpu_num_cores": "None", "tpu_metrics_debug": false, "debug": "[]", "dataloader_drop_last": false, "eval_steps": "None", "dataloader_num_workers": 0, "past_index": -1, "run_name": "/content/drive/MyDrive/Programming/hf-trainer/long-t5-tglobal-base-16384-booksum-V7.9-ft1-booksum", "disable_tqdm": false, "remove_unused_columns": true, "label_names": "None", "load_best_model_at_end": false, "metric_for_best_model": "None", "greater_is_better": "None", "ignore_data_skip": false, "sharded_ddp": "[]", "fsdp": "[]", "fsdp_min_num_params": 0, "deepspeed": "/content/ds_config_zero2.json", "label_smoothing_factor": 0.0, "optim": "adamw_hf", "adafactor": false, "group_by_length": false, "length_column_name": "length", "report_to": "['tensorboard']", "ddp_find_unused_parameters": "None", "ddp_bucket_cap_mb": "None", "dataloader_pin_memory": true, "skip_memory_metrics": true, "use_legacy_prediction_loop": false, "push_to_hub": true, "resume_from_checkpoint": "None", "hub_model_id": "long-t5-tglobal-base-16384-booksum-V7.9-ft1-booksum", "hub_strategy": "end", "hub_token": "<HUB_TOKEN>", "hub_private_repo": true, "gradient_checkpointing": true, "include_inputs_for_metrics": false, "fp16_backend": "auto", "push_to_hub_model_id": "None", "push_to_hub_organization": "None", "push_to_hub_token": "<PUSH_TO_HUB_TOKEN>", "_n_gpu": 1, "mp_parameters": "", "auto_find_batch_size": false, "full_determinism": false, "torchdynamo": "None", "ray_scope": "last", "sortish_sampler": false, "predict_with_generate": false, "generation_max_length": "None", "generation_num_beams": "None", "train_batch_size": 1, "eval_batch_size": 1, "configs_src": "long-t5-tglobal-base-16384-booksum-V7.9-ft1-booksum"}

long-t5-tglobal-base-16384-booksum-V9-ft1-booksum_training_metadata.json

@@ -0,0 +1 @@
+{"output_dir": "/content/drive/MyDrive/Programming/hf-trainer/long-t5-tglobal-base-16384-booksum-V9-ft1-booksum", "overwrite_output_dir": true, "do_train": false, "do_eval": false, "do_predict": false, "evaluation_strategy": "no", "prediction_loss_only": false, "per_device_train_batch_size": 1, "per_device_eval_batch_size": 1, "per_gpu_train_batch_size": "None", "per_gpu_eval_batch_size": "None", "gradient_accumulation_steps": 64, "eval_accumulation_steps": "None", "eval_delay": 0, "learning_rate": 0.0006, "weight_decay": 0.01, "adam_beta1": 0.9, "adam_beta2": 0.999, "adam_epsilon": 1e-08, "max_grad_norm": 0.3, "num_train_epochs": 2, "max_steps": -1, "lr_scheduler_type": "cosine", "warmup_ratio": 0.01, "warmup_steps": 0, "log_level": -1, "log_level_replica": -1, "log_on_each_node": true, "logging_dir": "/content/drive/MyDrive/Programming/hf-trainer/long-t5-tglobal-base-16384-booksum-V9-ft1-booksum/logs", "logging_strategy": "steps", "logging_first_step": false, "logging_steps": 2, "logging_nan_inf_filter": true, "save_strategy": "steps", "save_steps": 25, "save_total_limit": 1, "save_on_each_node": false, "no_cuda": false, "seed": 42, "data_seed": "None", "jit_mode_eval": false, "use_ipex": false, "bf16": false, "fp16": true, "fp16_opt_level": "O1", "half_precision_backend": "cuda_amp", "bf16_full_eval": false, "fp16_full_eval": false, "tf32": "None", "local_rank": 0, "xpu_backend": "None", "tpu_num_cores": "None", "tpu_metrics_debug": false, "debug": "[]", "dataloader_drop_last": false, "eval_steps": "None", "dataloader_num_workers": 0, "past_index": -1, "run_name": "/content/drive/MyDrive/Programming/hf-trainer/long-t5-tglobal-base-16384-booksum-V9-ft1-booksum", "disable_tqdm": false, "remove_unused_columns": true, "label_names": "None", "load_best_model_at_end": false, "metric_for_best_model": "None", "greater_is_better": "None", "ignore_data_skip": false, "sharded_ddp": "[]", "fsdp": "[]", "fsdp_min_num_params": 0, "deepspeed": "/content/ds_config_zero2.json", "label_smoothing_factor": 0.0, "optim": "adamw_hf", "adafactor": false, "group_by_length": false, "length_column_name": "length", "report_to": "['tensorboard']", "ddp_find_unused_parameters": "None", "ddp_bucket_cap_mb": "None", "dataloader_pin_memory": true, "skip_memory_metrics": true, "use_legacy_prediction_loop": false, "push_to_hub": true, "resume_from_checkpoint": "None", "hub_model_id": "long-t5-tglobal-base-16384-booksum-V9-ft1-booksum", "hub_strategy": "end", "hub_token": "<HUB_TOKEN>", "hub_private_repo": true, "gradient_checkpointing": true, "include_inputs_for_metrics": false, "fp16_backend": "auto", "push_to_hub_model_id": "None", "push_to_hub_organization": "None", "push_to_hub_token": "<PUSH_TO_HUB_TOKEN>", "_n_gpu": 1, "mp_parameters": "", "auto_find_batch_size": false, "full_determinism": false, "torchdynamo": "None", "ray_scope": "last", "sortish_sampler": false, "predict_with_generate": false, "generation_max_length": "None", "generation_num_beams": "None", "train_batch_size": 1, "eval_batch_size": 1, "configs_src": "long-t5-tglobal-base-16384-booksum-V9-ft1-booksum"}

model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:851dbab14fe12a66c8b5f8efa4a2135061db5ef54ef87b609afb3b27519c2b85
+size 992593930

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:134d0d187b6cde13e94f8ae3dca1aecdc26ecafe075410af4880e5ec9cf0c23d
+size 990446387

special_tokens_map.json

@@ -0,0 +1,107 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "eos_token": "</s>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}

spiece.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d60acb128cf7b7f2536e8f38a5b18a05535c9e14c7a355904270e15b0945ea86
+size 791656

tokenizer_config.json

@@ -0,0 +1,111 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "eos_token": "</s>",
+  "extra_ids": 100,
+  "name_or_path": "pszemraj/long-t5-tglobal-base-16384-booksum-V10",
+  "pad_token": "<pad>",
+  "special_tokens_map_file": null,
+  "tokenizer_class": "T5Tokenizer",
+  "unk_token": "<unk>"
+}

trainer_state.json

@@ -0,0 +1,1015 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 1.9971671388101981,
+  "global_step": 330,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.01,
+      "learning_rate": 0.0003,
+      "loss": 2.0736,
+      "step": 2
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 0.0006,
+      "loss": 2.0741,
+      "step": 4
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 0.0005999442811234004,
+      "loss": 2.098,
+      "step": 6
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 0.0005997771451908898,
+      "loss": 2.0627,
+      "step": 8
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 0.0005994986542866444,
+      "loss": 2.0418,
+      "step": 10
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 0.0005991089118586658,
+      "loss": 2.0717,
+      "step": 12
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 0.0005986080626803564,
+      "loss": 2.0504,
+      "step": 14
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 0.0005979962927967394,
+      "loss": 2.0314,
+      "step": 16
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 0.0005972738294553527,
+      "loss": 2.0568,
+      "step": 18
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 0.0005964409410218346,
+      "loss": 2.073,
+      "step": 20
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 0.0005954979368802377,
+      "loss": 2.0737,
+      "step": 22
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 0.0005944451673181043,
+      "loss": 2.0936,
+      "step": 24
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 0.0005932830233963502,
+      "loss": 2.0846,
+      "step": 26
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 0.0005920119368040003,
+      "loss": 2.1374,
+      "step": 28
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 0.0005906323796978346,
+      "loss": 2.0828,
+      "step": 30
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 0.0005891448645270008,
+      "loss": 2.0802,
+      "step": 32
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 0.0005875499438426604,
+      "loss": 2.147,
+      "step": 34
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 0.0005858482100927391,
+      "loss": 2.0904,
+      "step": 36
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 0.0005840402954018554,
+      "loss": 2.0823,
+      "step": 38
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 0.0005821268713365129,
+      "loss": 2.0792,
+      "step": 40
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 0.0005801086486556411,
+      "loss": 2.0716,
+      "step": 42
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 0.0005779863770465765,
+      "loss": 2.2061,
+      "step": 44
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 0.0005757608448465853,
+      "loss": 2.1735,
+      "step": 46
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 0.0005734328787500274,
+      "loss": 2.098,
+      "step": 48
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 0.000571003343501274,
+      "loss": 2.0977,
+      "step": 50
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 0.0005684731415734899,
+      "loss": 2.1076,
+      "step": 52
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 0.0005658432128334017,
+      "loss": 2.0476,
+      "step": 54
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 0.0005631145341921755,
+      "loss": 2.0728,
+      "step": 56
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 0.0005602881192425346,
+      "loss": 2.0477,
+      "step": 58
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 0.000557365017882251,
+      "loss": 2.1102,
+      "step": 60
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 0.0005543463159241515,
+      "loss": 2.0754,
+      "step": 62
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 0.0005512331346927817,
+      "loss": 2.1029,
+      "step": 64
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 0.0005480266306078807,
+      "loss": 1.9922,
+      "step": 66
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 0.0005447279947548182,
+      "loss": 2.1599,
+      "step": 68
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 0.0005413384524421545,
+      "loss": 2.1388,
+      "step": 70
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 0.0005378592627464883,
+      "loss": 2.131,
+      "step": 72
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 0.0005342917180447621,
+      "loss": 2.0658,
+      "step": 74
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 0.0005306371435341955,
+      "loss": 2.1437,
+      "step": 76
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 0.0005268968967400301,
+      "loss": 2.1145,
+      "step": 78
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 0.0005230723670112627,
+      "loss": 2.1263,
+      "step": 80
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 0.0005191649750045603,
+      "loss": 2.0513,
+      "step": 82
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 0.000515176172156543,
+      "loss": 2.0723,
+      "step": 84
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 0.0005111074401446355,
+      "loss": 2.0744,
+      "step": 86
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 0.0005069602903366834,
+      "loss": 2.131,
+      "step": 88
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 0.0005027362632295429,
+      "loss": 2.094,
+      "step": 90
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 0.0004984369278768482,
+      "loss": 2.0633,
+      "step": 92
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 0.0004940638813061723,
+      "loss": 2.1182,
+      "step": 94
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 0.0004896187479257971,
+      "loss": 2.1664,
+      "step": 96
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 0.0004851031789213108,
+      "loss": 2.1365,
+      "step": 98
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 0.0004805188516422613,
+      "loss": 2.1204,
+      "step": 100
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 0.00047586746897908803,
+      "loss": 2.0677,
+      "step": 102
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 0.00047115075873056876,
+      "loss": 2.1588,
+      "step": 104
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 0.0004663704729620119,
+      "loss": 2.0746,
+      "step": 106
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 0.0004615283873544366,
+      "loss": 2.126,
+      "step": 108
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 0.0004566263005449791,
+      "loss": 2.0786,
+      "step": 110
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 0.0004516660334587729,
+      "loss": 2.1019,
+      "step": 112
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 0.00044664942863254913,
+      "loss": 2.0605,
+      "step": 114
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 0.0004415783495302096,
+      "loss": 2.0879,
+      "step": 116
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 0.0004364546798506258,
+      "loss": 2.1516,
+      "step": 118
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 0.0004312803228279218,
+      "loss": 2.1287,
+      "step": 120
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 0.0004260572005245005,
+      "loss": 2.1054,
+      "step": 122
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 0.00042078725311707585,
+      "loss": 2.0905,
+      "step": 124
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 0.0004154724381759763,
+      "loss": 2.0875,
+      "step": 126
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 0.0004101147299379876,
+      "loss": 2.0948,
+      "step": 128
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 0.00040471611857300423,
+      "loss": 2.1521,
+      "step": 130
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 0.0003992786094447625,
+      "loss": 2.0645,
+      "step": 132
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 0.0003938042223659299,
+      "loss": 2.0953,
+      "step": 134
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 0.0003882949908478272,
+      "loss": 2.1046,
+      "step": 136
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 0.0003827529613450617,
+      "loss": 2.1814,
+      "step": 138
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 0.0003771801924953534,
+      "loss": 2.1147,
+      "step": 140
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 0.0003715787543548345,
+      "loss": 2.1226,
+      "step": 142
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 0.0003659507276291087,
+      "loss": 2.1093,
+      "step": 144
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 0.00036029820290035347,
+      "loss": 2.1038,
+      "step": 146
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 0.0003546232798507543,
+      "loss": 2.0581,
+      "step": 148
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 0.00034892806648255875,
+      "loss": 2.0076,
+      "step": 150
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 0.0003432146783350393,
+      "loss": 2.1017,
+      "step": 152
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 0.0003374852376986568,
+      "loss": 2.1353,
+      "step": 154
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 0.00033174187282671686,
+      "loss": 2.0836,
+      "step": 156
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 0.0003259867171448097,
+      "loss": 2.098,
+      "step": 158
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 0.00032022190845833035,
+      "loss": 2.1308,
+      "step": 160
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 0.0003144495881583712,
+      "loss": 2.1204,
+      "step": 162
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 0.00030867190042628177,
+      "loss": 2.0564,
+      "step": 164
+    },
+    {
+      "epoch": 1.01,
+      "learning_rate": 0.0003028909914371942,
+      "loss": 2.5573,
+      "step": 166
+    },
+    {
+      "epoch": 1.02,
+      "learning_rate": 0.0002971090085628058,
+      "loss": 1.9576,
+      "step": 168
+    },
+    {
+      "epoch": 1.03,
+      "learning_rate": 0.00029132809957371823,
+      "loss": 1.9274,
+      "step": 170
+    },
+    {
+      "epoch": 1.04,
+      "learning_rate": 0.0002855504118416288,
+      "loss": 1.9637,
+      "step": 172
+    },
+    {
+      "epoch": 1.05,
+      "learning_rate": 0.0002797780915416696,
+      "loss": 1.9729,
+      "step": 174
+    },
+    {
+      "epoch": 1.07,
+      "learning_rate": 0.0002740132828551904,
+      "loss": 2.0254,
+      "step": 176
+    },
+    {
+      "epoch": 1.08,
+      "learning_rate": 0.00026825812717328314,
+      "loss": 1.9427,
+      "step": 178
+    },
+    {
+      "epoch": 1.09,
+      "learning_rate": 0.00026251476230134313,
+      "loss": 1.97,
+      "step": 180
+    },
+    {
+      "epoch": 1.1,
+      "learning_rate": 0.00025678532166496077,
+      "loss": 2.042,
+      "step": 182
+    },
+    {
+      "epoch": 1.11,
+      "learning_rate": 0.00025107193351744115,
+      "loss": 2.0033,
+      "step": 184
+    },
+    {
+      "epoch": 1.13,
+      "learning_rate": 0.0002453767201492457,
+      "loss": 1.9466,
+      "step": 186
+    },
+    {
+      "epoch": 1.14,
+      "learning_rate": 0.00023970179709964656,
+      "loss": 2.0294,
+      "step": 188
+    },
+    {
+      "epoch": 1.15,
+      "learning_rate": 0.0002340492723708912,
+      "loss": 2.0002,
+      "step": 190
+    },
+    {
+      "epoch": 1.16,
+      "learning_rate": 0.00022842124564516548,
+      "loss": 1.9686,
+      "step": 192
+    },
+    {
+      "epoch": 1.18,
+      "learning_rate": 0.0002228198075046467,
+      "loss": 2.0857,
+      "step": 194
+    },
+    {
+      "epoch": 1.19,
+      "learning_rate": 0.00021724703865493827,
+      "loss": 2.0111,
+      "step": 196
+    },
+    {
+      "epoch": 1.2,
+      "learning_rate": 0.00021170500915217283,
+      "loss": 2.0058,
+      "step": 198
+    },
+    {
+      "epoch": 1.21,
+      "learning_rate": 0.00020619577763407015,
+      "loss": 1.9725,
+      "step": 200
+    },
+    {
+      "epoch": 1.22,
+      "learning_rate": 0.0002007213905552375,
+      "loss": 2.0542,
+      "step": 202
+    },
+    {
+      "epoch": 1.24,
+      "learning_rate": 0.0001952838814269958,
+      "loss": 2.0265,
+      "step": 204
+    },
+    {
+      "epoch": 1.25,
+      "learning_rate": 0.00018988527006201237,
+      "loss": 2.1143,
+      "step": 206
+    },
+    {
+      "epoch": 1.26,
+      "learning_rate": 0.00018452756182402364,
+      "loss": 1.96,
+      "step": 208
+    },
+    {
+      "epoch": 1.27,
+      "learning_rate": 0.00017921274688292415,
+      "loss": 2.0181,
+      "step": 210
+    },
+    {
+      "epoch": 1.28,
+      "learning_rate": 0.00017394279947549948,
+      "loss": 1.9909,
+      "step": 212
+    },
+    {
+      "epoch": 1.3,
+      "learning_rate": 0.00016871967717207824,
+      "loss": 2.0021,
+      "step": 214
+    },
+    {
+      "epoch": 1.31,
+      "learning_rate": 0.00016354532014937418,
+      "loss": 1.9598,
+      "step": 216
+    },
+    {
+      "epoch": 1.32,
+      "learning_rate": 0.00015842165046979042,
+      "loss": 1.9433,
+      "step": 218
+    },
+    {
+      "epoch": 1.33,
+      "learning_rate": 0.0001533505713674509,
+      "loss": 2.0222,
+      "step": 220
+    },
+    {
+      "epoch": 1.34,
+      "learning_rate": 0.0001483339665412271,
+      "loss": 1.9898,
+      "step": 222
+    },
+    {
+      "epoch": 1.36,
+      "learning_rate": 0.00014337369945502084,
+      "loss": 2.0555,
+      "step": 224
+    },
+    {
+      "epoch": 1.37,
+      "learning_rate": 0.00013847161264556339,
+      "loss": 2.0082,
+      "step": 226
+    },
+    {
+      "epoch": 1.38,
+      "learning_rate": 0.000133629527037988,
+      "loss": 2.0081,
+      "step": 228
+    },
+    {
+      "epoch": 1.39,
+      "learning_rate": 0.00012884924126943122,
+      "loss": 2.0028,
+      "step": 230
+    },
+    {
+      "epoch": 1.4,
+      "learning_rate": 0.00012413253102091197,
+      "loss": 2.0535,
+      "step": 232
+    },
+    {
+      "epoch": 1.42,
+      "learning_rate": 0.00011948114835773868,
+      "loss": 1.9512,
+      "step": 234
+    },
+    {
+      "epoch": 1.43,
+      "learning_rate": 0.00011489682107868918,
+      "loss": 1.9141,
+      "step": 236
+    },
+    {
+      "epoch": 1.44,
+      "learning_rate": 0.00011038125207420298,
+      "loss": 2.0705,
+      "step": 238
+    },
+    {
+      "epoch": 1.45,
+      "learning_rate": 0.00010593611869382759,
+      "loss": 1.9869,
+      "step": 240
+    },
+    {
+      "epoch": 1.47,
+      "learning_rate": 0.0001015630721231517,
+      "loss": 1.9448,
+      "step": 242
+    },
+    {
+      "epoch": 1.48,
+      "learning_rate": 9.7263736770457e-05,
+      "loss": 1.9565,
+      "step": 244
+    },
+    {
+      "epoch": 1.49,
+      "learning_rate": 9.303970966331645e-05,
+      "loss": 1.9925,
+      "step": 246
+    },
+    {
+      "epoch": 1.5,
+      "learning_rate": 8.88925598553645e-05,
+      "loss": 1.995,
+      "step": 248
+    },
+    {
+      "epoch": 1.51,
+      "learning_rate": 8.482382784345695e-05,
+      "loss": 1.9627,
+      "step": 250
+    },
+    {
+      "epoch": 1.53,
+      "learning_rate": 8.083502499543967e-05,
+      "loss": 2.0096,
+      "step": 252
+    },
+    {
+      "epoch": 1.54,
+      "learning_rate": 7.692763298873725e-05,
+      "loss": 2.0731,
+      "step": 254
+    },
+    {
+      "epoch": 1.55,
+      "learning_rate": 7.310310325996986e-05,
+      "loss": 2.0134,
+      "step": 256
+    },
+    {
+      "epoch": 1.56,
+      "learning_rate": 6.936285646580441e-05,
+      "loss": 2.0516,
+      "step": 258
+    },
+    {
+      "epoch": 1.57,
+      "learning_rate": 6.570828195523786e-05,
+      "loss": 2.0061,
+      "step": 260
+    },
+    {
+      "epoch": 1.59,
+      "learning_rate": 6.214073725351162e-05,
+      "loss": 1.9709,
+      "step": 262
+    },
+    {
+      "epoch": 1.6,
+      "learning_rate": 5.8661547557845534e-05,
+      "loss": 2.03,
+      "step": 264
+    },
+    {
+      "epoch": 1.61,
+      "learning_rate": 5.5272005245181705e-05,
+      "loss": 1.9751,
+      "step": 266
+    },
+    {
+      "epoch": 1.62,
+      "learning_rate": 5.197336939211925e-05,
+      "loss": 2.0363,
+      "step": 268
+    },
+    {
+      "epoch": 1.63,
+      "learning_rate": 4.87668653072184e-05,
+      "loss": 2.0713,
+      "step": 270
+    },
+    {
+      "epoch": 1.65,
+      "learning_rate": 4.565368407584855e-05,
+      "loss": 2.021,
+      "step": 272
+    },
+    {
+      "epoch": 1.66,
+      "learning_rate": 4.2634982117748854e-05,
+      "loss": 1.9747,
+      "step": 274
+    },
+    {
+      "epoch": 1.67,
+      "learning_rate": 3.971188075746535e-05,
+      "loss": 2.0548,
+      "step": 276
+    },
+    {
+      "epoch": 1.68,
+      "learning_rate": 3.688546580782448e-05,
+      "loss": 1.9953,
+      "step": 278
+    },
+    {
+      "epoch": 1.69,
+      "learning_rate": 3.415678716659831e-05,
+      "loss": 1.9769,
+      "step": 280
+    },
+    {
+      "epoch": 1.71,
+      "learning_rate": 3.152685842651004e-05,
+      "loss": 2.0558,
+      "step": 282
+    },
+    {
+      "epoch": 1.72,
+      "learning_rate": 2.899665649872589e-05,
+      "loss": 2.047,
+      "step": 284
+    },
+    {
+      "epoch": 1.73,
+      "learning_rate": 2.6567121249972544e-05,
+      "loss": 2.0231,
+      "step": 286
+    },
+    {
+      "epoch": 1.74,
+      "learning_rate": 2.423915515341467e-05,
+      "loss": 2.0504,
+      "step": 288
+    },
+    {
+      "epoch": 1.76,
+      "learning_rate": 2.2013622953423405e-05,
+      "loss": 2.0075,
+      "step": 290
+    },
+    {
+      "epoch": 1.77,
+      "learning_rate": 1.9891351344358853e-05,
+      "loss": 2.0219,
+      "step": 292
+    },
+    {
+      "epoch": 1.78,
+      "learning_rate": 1.7873128663487047e-05,
+      "loss": 1.9923,
+      "step": 294
+    },
+    {
+      "epoch": 1.79,
+      "learning_rate": 1.5959704598144628e-05,
+      "loss": 2.0081,
+      "step": 296
+    },
+    {
+      "epoch": 1.8,
+      "learning_rate": 1.4151789907260846e-05,
+      "loss": 1.9971,
+      "step": 298
+    },
+    {
+      "epoch": 1.82,
+      "learning_rate": 1.2450056157339439e-05,
+      "loss": 1.9923,
+      "step": 300
+    },
+    {
+      "epoch": 1.83,
+      "learning_rate": 1.0855135472999155e-05,
+      "loss": 2.003,
+      "step": 302
+    },
+    {
+      "epoch": 1.84,
+      "learning_rate": 9.36762030216538e-06,
+      "loss": 2.0014,
+      "step": 304
+    },
+    {
+      "epoch": 1.85,
+      "learning_rate": 7.988063195999639e-06,
+      "loss": 2.0675,
+      "step": 306
+    },
+    {
+      "epoch": 1.86,
+      "learning_rate": 6.716976603649782e-06,
+      "loss": 1.9773,
+      "step": 308
+    },
+    {
+      "epoch": 1.88,
+      "learning_rate": 5.554832681895582e-06,
+      "loss": 1.997,
+      "step": 310
+    },
+    {
+      "epoch": 1.89,
+      "learning_rate": 4.5020631197623e-06,
+      "loss": 2.0323,
+      "step": 312
+    },
+    {
+      "epoch": 1.9,
+      "learning_rate": 3.559058978165319e-06,
+      "loss": 2.0221,
+      "step": 314
+    },
+    {
+      "epoch": 1.91,
+      "learning_rate": 2.7261705446473103e-06,
+      "loss": 1.9734,
+      "step": 316
+    },
+    {
+      "epoch": 1.92,
+      "learning_rate": 2.003707203260585e-06,
+      "loss": 1.9529,
+      "step": 318
+    },
+    {
+      "epoch": 1.94,
+      "learning_rate": 1.3919373196436213e-06,
+      "loss": 1.982,
+      "step": 320
+    },
+    {
+      "epoch": 1.95,
+      "learning_rate": 8.910881413340398e-07,
+      "loss": 2.0501,
+      "step": 322
+    },
+    {
+      "epoch": 1.96,
+      "learning_rate": 5.013457133556276e-07,
+      "loss": 2.0041,
+      "step": 324
+    },
+    {
+      "epoch": 1.97,
+      "learning_rate": 2.2285480911008457e-07,
+      "loss": 1.983,
+      "step": 326
+    },
+    {
+      "epoch": 1.99,
+      "learning_rate": 5.5718876599541995e-08,
+      "loss": 2.0002,
+      "step": 328
+    },
+    {
+      "epoch": 2.0,
+      "learning_rate": 0.0,
+      "loss": 2.0595,
+      "step": 330
+    },
+    {
+      "epoch": 2.0,
+      "step": 330,
+      "total_flos": 4.634629372945367e+17,
+      "train_loss": 2.053801321260857,
+      "train_runtime": 80029.5559,
+      "train_samples_per_second": 0.265,
+      "train_steps_per_second": 0.004
+    }
+  ],
+  "max_steps": 330,
+  "num_train_epochs": 2,
+  "total_flos": 4.634629372945367e+17,
+  "trial_name": null,
+  "trial_params": null
+}

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7590081405196835a42cca00f60db01b8ca901bae8a197f0becb2fcfa0999a94
+size 4527

zero_to_fp32.py

@@ -0,0 +1,484 @@
+#!/usr/bin/env python
+
+# This script extracts fp32 consolidated weights from a zero 2 and 3 DeepSpeed checkpoints. It gets
+# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
+# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
+# application.
+#
+# example: python zero_to_fp32.py . pytorch_model.bin
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+from collections import OrderedDict
+
+# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
+# DeepSpeed data structures it has to be available in the current python environment.
+import deepspeed
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION,
+                                            OPTIMIZER_STATE_DICT,
+                                            PARAM_SHAPES,
+                                            SINGLE_PARTITION_OF_FP32_GROUPS,
+                                            FP32_FLAT_GROUPS,
+                                            ZERO_STAGE,
+                                            PARTITION_COUNT,
+                                            PARAM_SHAPES,
+                                            BUFFER_NAMES)
+
+debug = 0
+
+# load to cpu
+device = torch.device('cpu')
+
+
+def atoi(text):
+    return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+    '''
+    alist.sort(key=natural_keys) sorts in human order
+    http://nedbatchelder.com/blog/200712/human_sorting.html
+    (See Toothy's implementation in the comments)
+    '''
+    return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
+def get_model_state_file(checkpoint_dir, zero_stage):
+    if not os.path.isdir(checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
+
+    # there should be only one file
+    if zero_stage == 2:
+        file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
+    elif zero_stage == 3:
+        file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
+
+    if not os.path.exists(file):
+        raise FileNotFoundError(f"can't find model states file at '{file}'")
+
+    return file
+
+
+def get_optim_files(checkpoint_dir):
+    # XXX: need to test that this simple glob rule works for multi-node setup too
+    optim_files = sorted(glob.glob(os.path.join(checkpoint_dir,
+                                                "*_optim_states.pt")),
+                         key=natural_keys)
+
+    if len(optim_files) == 0:
+        raise FileNotFoundError(
+            f"can't find '*_optim_states.pt' files in directory '{checkpoint_dir}'")
+
+    return optim_files
+
+
+def parse_model_state(file):
+    state_dict = torch.load(file, map_location=device)
+
+    if BUFFER_NAMES not in state_dict:
+        raise ValueError(f"{file} is not a model state checkpoint")
+    buffer_names = state_dict[BUFFER_NAMES]
+    if debug:
+        print("Found buffers:", buffer_names)
+
+    # recover just the buffers while restoring them to fp32 if they were saved in fp16
+    buffers = {
+        k: v.float()
+        for k,
+        v in state_dict["module"].items() if k in buffer_names
+    }
+    param_shapes = state_dict[PARAM_SHAPES]
+
+    ds_version = state_dict.get(DS_VERSION, None)
+
+    return buffers, param_shapes, ds_version
+
+
+def parse_optim_states(files, ds_checkpoint_dir):
+
+    total_files = len(files)
+    state_dicts = []
+    for f in files:
+        state_dicts.append(torch.load(f, map_location=device))
+
+    if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]:
+        raise ValueError(f"{files[0]} is not a zero checkpoint")
+    zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
+    world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
+
+    # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
+    # parameters can be different from data parallelism for non-expert parameters. So we can just
+    # use the max of the partition_count to get the dp world_size.
+
+    if type(world_size) is list:
+        world_size = max(world_size)
+
+    if world_size != total_files:
+        raise ValueError(
+            f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
+            "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
+        )
+
+    # the groups are named differently in each stage
+    if zero_stage == 2:
+        fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
+    elif zero_stage == 3:
+        fp32_groups_key = FP32_FLAT_GROUPS
+    else:
+        raise ValueError(f"unknown zero stage {zero_stage}")
+
+    if zero_stage == 2:
+        fp32_flat_groups = [
+            state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key]
+            for i in range(len(state_dicts))
+        ]
+    elif zero_stage == 3:
+        # if there is more than one param group, there will be multiple flattened tensors - one
+        # flattened tensor per group - for simplicity merge them into a single tensor
+        #
+        # XXX: could make the script more memory efficient for when there are multiple groups - it
+        # will require matching the sub-lists of param_shapes for each param group flattened tensor
+
+        fp32_flat_groups = [
+            torch.cat(state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key],
+                      0) for i in range(len(state_dicts))
+        ]
+
+    return zero_stage, world_size, fp32_flat_groups
+
+
+def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir):
+    """
+    Returns fp32 state_dict reconstructed from ds checkpoint
+
+    Args:
+        - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
+
+    """
+    print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
+
+    optim_files = get_optim_files(ds_checkpoint_dir)
+    zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
+    print(
+        f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
+
+    model_file = get_model_state_file(ds_checkpoint_dir, zero_stage)
+    buffers, param_shapes, ds_version = parse_model_state(model_file)
+    print(f'Parsing checkpoint created by deepspeed=={ds_version}')
+
+    if zero_stage == 2:
+        return _get_fp32_state_dict_from_zero2_checkpoint(world_size,
+                                                          param_shapes,
+                                                          fp32_flat_groups,
+                                                          buffers)
+    elif zero_stage == 3:
+        return _get_fp32_state_dict_from_zero3_checkpoint(world_size,
+                                                          param_shapes,
+                                                          fp32_flat_groups,
+                                                          buffers)
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size,
+                                               param_shapes,
+                                               fp32_flat_groups,
+                                               buffers):
+
+    # Reconstruction protocol:
+    #
+    # XXX: document this
+
+    if debug:
+        for i in range(world_size):
+            for j in range(len(fp32_flat_groups[0])):
+                print(
+                    f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
+
+    # XXX: memory usage doubles here (zero2)
+    num_param_groups = len(fp32_flat_groups[0])
+    merged_single_partition_of_fp32_groups = []
+    for i in range(num_param_groups):
+        merged_partitions = [sd[i] for sd in fp32_flat_groups]
+        full_single_fp32_vector = torch.cat(merged_partitions, 0)
+        merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
+    avail_numel = sum([
+        full_single_fp32_vector.numel()
+        for full_single_fp32_vector in merged_single_partition_of_fp32_groups
+    ])
+
+    if debug:
+        wanted_params = sum([len(shapes) for shapes in param_shapes])
+        wanted_numel = sum(
+            [sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
+        # not asserting if there is a mismatch due to possible padding
+        print(f"Have {avail_numel} numels to process.")
+        print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+    state_dict = OrderedDict()
+
+    # buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    total_numel = 0
+    total_params = 0
+    for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
+        offset = 0
+        avail_numel = full_single_fp32_vector.numel()
+        for name, shape in shapes.items():
+
+            unpartitioned_numel = shape.numel()
+            total_numel += unpartitioned_numel
+            total_params += 1
+
+            if debug:
+                print(
+                    f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} "
+                )
+            state_dict[name] = full_single_fp32_vector.narrow(
+                0,
+                offset,
+                unpartitioned_numel).view(shape)
+            offset += unpartitioned_numel
+
+        # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
+        # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
+        # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
+        # live optimizer object, so we are checking that the numbers are within the right range
+        align_to = 2 * world_size
+
+        def zero2_align(x):
+            return align_to * math.ceil(x / align_to)
+
+        if debug:
+            print(f"original offset={offset}, avail_numel={avail_numel}")
+
+        offset = zero2_align(offset)
+        avail_numel = zero2_align(avail_numel)
+
+        if debug:
+            print(f"aligned  offset={offset}, avail_numel={avail_numel}")
+
+        # Sanity check
+        if offset != avail_numel:
+            raise ValueError(
+                f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(
+        f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements"
+    )
+
+    return state_dict
+
+
+def zero3_partitioned_param_info(unpartitioned_numel, world_size):
+    remainder = unpartitioned_numel % world_size
+    padding_numel = (world_size - remainder) if remainder else 0
+    partitioned_numel = math.ceil(unpartitioned_numel / world_size)
+    return partitioned_numel, padding_numel
+
+
+def _get_fp32_state_dict_from_zero3_checkpoint(world_size,
+                                               param_shapes,
+                                               fp32_flat_groups,
+                                               buffers):
+
+    # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
+    # param, re-consolidating each param, while dealing with padding if any
+
+    avail_numel = fp32_flat_groups[0].numel() * world_size
+    # merge list of dicts, preserving order
+    param_shapes = {k: v for d in param_shapes for k, v in d.items()}
+
+    if debug:
+        for i in range(world_size):
+            print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
+
+        wanted_params = len(param_shapes)
+        wanted_numel = sum(shape.numel() for shape in param_shapes.values())
+        # not asserting if there is a mismatch due to possible padding
+        print(f"Have {avail_numel} numels to process.")
+        print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+    state_dict = OrderedDict()
+
+    # buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    offset = 0
+    total_numel = 0
+    total_params = 0
+    for name, shape in param_shapes.items():
+
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+        total_params += 1
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"{total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+        # XXX: memory usage doubles here
+        state_dict[name] = torch.cat(
+            tuple(fp32_flat_groups[i].narrow(0,
+                                             offset,
+                                             partitioned_numel)
+                  for i in range(world_size)),
+            0).narrow(0,
+                      0,
+                      unpartitioned_numel).view(shape)
+        offset += partitioned_numel
+
+    offset *= world_size
+
+    # Sanity check
+    if offset != avail_numel:
+        raise ValueError(
+            f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(
+        f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements"
+    )
+
+    return state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag=None):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
+    ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
+    via a model hub.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
+
+    Returns:
+        - pytorch ``state_dict``
+
+    Note: this approach may not work if your application doesn't have sufficient free CPU memory and
+    you may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+    the checkpoint.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        # do the training and checkpoint saving
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
+        model = model.cpu() # move to cpu
+        model.load_state_dict(state_dict)
+        # submit to model hub or save the model to share with others
+
+    In this example the ``model`` will no longer be usable in the deepspeed context of the same
+    application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
+
+    """
+    if tag is None:
+        latest_path = os.path.join(checkpoint_dir, 'latest')
+        if os.path.isfile(latest_path):
+            with open(latest_path, 'r') as fd:
+                tag = fd.read().strip()
+        else:
+            raise ValueError(f"Unable to find 'latest' file at {latest_path}")
+
+    ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
+
+    if not os.path.isdir(ds_checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
+
+    return _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, output_file, tag=None):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
+    loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``output_file``: path to the pytorch fp32 state_dict output file (e.g. path/pytorch_model.bin)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+    """
+
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+    print(f"Saving fp32 state dict to {output_file}")
+    torch.save(state_dict, output_file)
+
+
+def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
+    """
+    1. Put the provided model to cpu
+    2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
+    3. Load it into the provided model
+
+    Args:
+        - ``model``: the model object to update
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+
+    Returns:
+        - ``model`: modified model
+
+    Make sure you have plenty of CPU memory available before you call this function. If you don't
+    have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
+    conveniently placed for you in the checkpoint folder.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+        model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
+        # submit to model hub or save the model to share with others
+
+    Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
+    of the same application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    """
+    logger.info(f"Extracting fp32 weights")
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+
+    logger.info(f"Overwriting model with fp32 weights")
+    model = model.cpu()
+    model.load_state_dict(state_dict, strict=False)
+
+    return model
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "checkpoint_dir",
+        type=str,
+        help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
+    parser.add_argument(
+        "output_file",
+        type=str,
+        help=
+        "path to the pytorch fp32 state_dict output file (e.g. path/checkpoint-12/pytorch_model.bin)"
+    )
+    parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
+    args = parser.parse_args()
+
+    debug = args.debug
+
+    convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir, args.output_file)