Package: lightning-transformers
Deprecation notice 🔒This repository has been archived (read-only) on Nov 21, 2022. Thanks to everyone who contributed tolightning-transformers, we feel it's time to move on.:hugs: Transformers canalready be easily trained using the Lightning :zap: Trainer. Here's a recent example from the community:https://sachinruk.github.io/blog/deep-learning/2022/11/07/t5-for-grammar-correction.html. Note that there areno limitations or workarounds, things just work out of the box.Thelightning-transformersrepo explored the possibility to provide task-specific modules and pre-baked defaults, at the cost of introducing extra abstractions. In the spirit of keeping ourselves focused, these abstractions are not something we wish to continue supporting.If you likedlightning-transformersand want to continue developing it in the future, feel free to fork the repo and choose another name for the project.Flexible components pairing :hugs: Transformers withPytorch Lightning:zap:Docs•CommunityInstallationpipinstalllightning-transformersFrom Sourcegitclonehttps://github.com/PyTorchLightning/lightning-transformers.gitcdlightning-transformers
pipinstall.What is Lightning-TransformersLightning Transformers providesLightningModules,LightningDataModulesandStrategiesto use :hugs: Transformers with thePyTorch Lightning Trainer.Quick RecipesTrainbert-base-casedon theCARERemotion dataset using the Text Classification task.importpytorch_lightningasplfromtransformersimportAutoTokenizerfromlightning_transformers.task.nlp.text_classificationimport(TextClassificationDataModule,TextClassificationTransformer,)tokenizer=AutoTokenizer.from_pretrained(pretrained_model_name_or_path="bert-base-cased")dm=TextClassificationDataModule(batch_size=1,dataset_name="emotion",max_length=512,tokenizer=tokenizer,)model=TextClassificationTransformer(pretrained_model_name_or_path="bert-base-cased",num_labels=dm.num_classes)trainer=pl.Trainer(accelerator="auto",devices="auto",max_epochs=1)trainer.fit(model,dm)Train a pre-trainedmt5-basebackbone on theWMT16dataset using the Translation task.importpytorch_lightningasplfromtransformersimportAutoTokenizerfromlightning_transformers.task.nlp.translationimport(TranslationTransformer,WMT16TranslationDataModule,)tokenizer=AutoTokenizer.from_pretrained(pretrained_model_name_or_path="google/mt5-base")model=TranslationTransformer(pretrained_model_name_or_path="google/mt5-base",n_gram=4,smooth=False,val_target_max_length=142,num_beams=None,compute_generate_metrics=True,)dm=WMT16TranslationDataModule(# WMT translation datasets: ['cs-en', 'de-en', 'fi-en', 'ro-en', 'ru-en', 'tr-en']dataset_config_name="ro-en",source_language="en",target_language="ro",max_source_length=128,max_target_length=128,padding="max_length",tokenizer=tokenizer,)trainer=pl.Trainer(accelerator="auto",devices="auto",max_epochs=1)trainer.fit(model,dm)Lightning Transformers supports a bunch of :hugs: tasks and datasets. See thedocumentation.Billion Parameter Model SupportBig Model InferenceIt's really easy to enable large model support for the pre-built LightningModule :hugs: tasks.Below is an example to enable automatic model partitioning (across CPU/GPU and even leveraging disk space) to run text generation using a 6B parameter model.importtorchfromaccelerateimportinit_empty_weightsfromtransformersimportAutoTokenizerfromlightning_transformers.task.nlp.language_modelingimport(LanguageModelingTransformer,)withinit_empty_weights():model=LanguageModelingTransformer(pretrained_model_name_or_path="EleutherAI/gpt-j-6B",tokenizer=AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B"),low_cpu_mem_usage=True,device_map="auto",# automatically partitions the model based on the available hardware.)output=model.generate("Hello, my name is",device=torch.device("cuda"))print(model.tokenizer.decode(output[0].tolist()))For more information seeBig Transformers Model Inference.Big Model Training with DeepSpeedBelow is an example of how you can train a 6B parameter transformer model using Lightning Transformers and DeepSpeed.importpytorch_lightningasplfromtransformersimportAutoTokenizerfromlightning_transformers.task.nlp.language_modelingimport(LanguageModelingDataModule,LanguageModelingTransformer,)tokenizer=AutoTokenizer.from_pretrained(pretrained_model_name_or_path="gpt2")model=LanguageModelingTransformer(pretrained_model_name_or_path="EleutherAI/gpt-j-6B",tokenizer=AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B"),deepspeed_sharding=True,# defer initialization of the model to shard/load pre-train weights)dm=LanguageModelingDataModule(batch_size=1,dataset_name="wikitext",dataset_config_name="wikitext-2-raw-v1",tokenizer=tokenizer,)trainer=pl.Trainer(accelerator="gpu",devices="auto",strategy="deepspeed_stage_3",precision=16,max_epochs=1,)trainer.fit(model,dm)For more information seeDeepSpeed Training with Big Transformers Modelsor theModel Parallelismdocumentation.ContributePull requests are welcome. For major changes, please open an issue first to discuss what you would like to change.Please make sure to update tests as appropriate.CommunityFor help or questions, join our huge community onSlack!