Diffusion-based hierarchical language modeling.

Dependencies

Please follow the instructions in genslm to setup environment. This is particularly important if you plan to use DeepSpeed for distributed training.

Next, install this directory by

pip install -e .

Training with DeepSpeed Zero Stage 2

For foundation models with fewer than (including) 2.5B parameters, we can train the model using Zero Stage 2:

``` export NODES=10 export GPUSPERNODE=4 export MASTERADDR=x3006c0s13b1n0.hsn.cm.polaris.alcf.anl.gov export LR=1e-4 export EPOCHS=20 export TRAINBATCHSIZE=2 export ACCUMULATION=1 export EVALBATCHSIZE=1 export SAVETOTALLIMIT=5 export SAVEFOLDER=2.5B${NODES}nodesdeepspeeddiffusionsepcheckpoints${LR} export TRAINFILE=data/sampletrain.txt export TESTFILE=data/sampleval.txt export CLMODEL=/lus/eagle/projects/CVD-Mol-AI/yuntian/genomenewnaive/encoder93810/runl0.001b32/checkpoints export MODEL=EleutherAI/gpt-neox-20b # doesn't matter, will be ignored deepspeed --numgpus=${GPUSPERNODE} --numnodes=${NODES} --masteraddr=${MASTERADDR} --hostfile=hostfile --masterport=54321 examples/pytorch/language-modeling/runclmgenslm2.5B.py \ --perdevicetrainbatchsize=${TRAINBATCHSIZE} \ --deepspeed=deepspeedconfigs/zero2.json \ --perdeviceevalbatchsize=${EVALBATCHSIZE} \ --gradientaccumulationsteps=${ACCUMULATION} \ --outputdir=${SAVEFOLDER} \ --modeltype=${MODEL} \ --modelnameorpath=${MODEL} \ --dotrain \ --doeval \ --trainfile=${TRAINFILE} \ --validationfile=${TESTFILE} --overwriteoutputdir --savetotallimit=${SAVETOTALLIMIT} \ --learningrate=${LR} --numtrainepochs=${EPOCHS} --loadbestmodelatend=True \ --evaluationstrategy=epoch --savestrategy=epoch \ --clmodelnameorpath=${CLMODEL} \ --latentdim=32 \ --blocksize 1024 --fp16 --predictionloss_only

```

Training with DeepSpeed Zero Stage 3

``` export NODES=10 export GPUSPERNODE=4 export MASTERADDR=x3006c0s13b1n0.hsn.cm.polaris.alcf.anl.gov export LR=1e-4 export EPOCHS=20 export TRAINBATCHSIZE=2 export ACCUMULATION=1 export EVALBATCHSIZE=1 export SAVETOTALLIMIT=5 export SAVEFOLDER=2.5B${NODES}nodesdeepspeeddiffusionsepcheckpoints${LR} export TRAINFILE=data/sampletrain.txt export TESTFILE=data/sampleval.txt export CLMODEL=/lus/eagle/projects/CVD-Mol-AI/yuntian/genomenewnaive/encoder93810/runl0.001b32/checkpoints export MODEL=EleutherAI/gpt-neox-20b # doesn't matter, will be ignored deepspeed --numgpus=${GPUSPERNODE} --numnodes=${NODES} --masteraddr=${MASTERADDR} --hostfile=hostfile --masterport=54321 examples/pytorch/language-modeling/runclmgenslm25B.py \ --perdevicetrainbatchsize=${TRAINBATCHSIZE} \ --deepspeed=deepspeedconfigs/zero3.json \ --perdeviceevalbatchsize=${EVALBATCHSIZE} \ --gradientaccumulationsteps=${ACCUMULATION} \ --outputdir=${SAVEFOLDER} \ --modeltype=${MODEL} \ --modelnameorpath=${MODEL} \ --dotrain \ --doeval \ --trainfile=${TRAINFILE} \ --validationfile=${TESTFILE} --overwriteoutputdir --savetotallimit=${SAVETOTALLIMIT} \ --learningrate=${LR} --numtrainepochs=${EPOCHS} --loadbestmodelatend=True \ --evaluationstrategy=epoch --savestrategy=epoch \ --clmodelnameorpath=${CLMODEL} \ --latentdim=32 \ --blocksize 1024 --fp16 --predictionloss_only

```

Generate

To generate, run

CUDA_VISIBLE_DEVICES=0 python examples/pytorch/language-modeling/generate_genslm_2.5B.py

Citations

If you use our models in your research, please cite this paper:

@article{zvyagin2022genslms, title={GenSLMs: Genome-scale language models reveal SARS-CoV-2 evolutionary dynamics.}, author={Zvyagin, Max T and Brace, Alexander and Hippe, Kyle and Deng, Yuntian and Zhang, Bin and Bohorquez, Cindy Orozco and Clyde, Austin and Kale, Bharat and Perez-Rivera, Danilo and Ma, Heng and others}, journal={bioRxiv}, year={2022}, publisher={Cold Spring Harbor Laboratory} }