novarobot
/
RWKV-LM
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

no message

Browse Source
main
BlinkDL 3 years ago
parent f46cf49852
commit 8cced0383e
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File

4
RWKV-v4/train.py
Unescape Escape View File

@ -67,8 +67,8 @@ os.environ['RWKV_FLOAT_MODE'] = 'bf16'
os.environ['RWKV_DEEPSPEED'] = '1' # Use DeepSpeed? 0 = False, 1 = True os.environ['RWKV_DEEPSPEED'] = '1' # Use DeepSpeed? 0 = False, 1 = True
if int(os.environ['RWKV_NUM_GPUS']) == 1: # turn off DeepSpeed for 1 GPU training if int(os.environ['RWKV_NUM_GPUS']) == 1: # Usually you don't need DeepSpeed for 1 GPU training.
os.environ['RWKV_DEEPSPEED'] = '0' os.environ['RWKV_DEEPSPEED'] = '0' # However, sometimes DeepSpeed saves VRAM even for 1 GPU training. So you shall try it.
os.environ['USE_WANDB'] = '0' # wandb logging. 0 = False, 1 = True os.environ['USE_WANDB'] = '0' # wandb logging. 0 = False, 1 = True

Write Preview
Loading…
Cancel
Save