AllenNLP Deepspeed Config

config.json

{
    "dataset_reader": {
        "type": "sst_tokens",
        "granularity": "2-class",
        "max_instances": 1000,
        "token_indexers": {
            "tokens": {
                "type": "pretrained_transformer",
                "model_name": "roberta-base"
            }
        },
        "tokenizer": {
            "type": "pretrained_transformer",
            "model_name": "roberta-base"
        }
    },
    "model": {
        "type": "basic_classifier",
        "seq2vec_encoder": {
            "type": "bert_pooler",
            "dropout": 0.1,
            "pretrained_model": "roberta-base"
        },
        "text_field_embedder": {
            "token_embedders": {
                "tokens": {
                    "type": "pretrained_transformer",
                    "model_name": "roberta-base"
                }
            }
        }
    },
    "train_data_path": "https://allennlp.s3.amazonaws.com/datasets/sst/train.txt",
    "trainer": {
        "type": "deepspeed",
        "deepspeed_config": {
            "fp16": {
                "enabled": true
            },
            "optimizer": {
                "type": "Adam",
                "params": {
                    "lr": 1e-05
                }
            },
            "zero_optimization": {
                "stage": 2
            }
        },
        "num_epochs": 5,
        "optimizer": "adam"
    },
    "data_loader": {
        "batch_size": 64,
        "shuffle": true
    },
    "distributed": {
        "cuda_devices": [
            0,
            1
        ]
    }
}