Set replace_with_kernel_inject=True in BERT fill-mask example (deepspeedai#225)

lekurile · web-flow · commit d5eb762eec7c · 2022-12-01T13:49:02.000-08:00
This PR sets replace_with_kernel_inject=True in the BERT fill-mask inference example.
diff --git a/inference/huggingface/fill-mask/test-bert.py b/inference/huggingface/fill-mask/test-bert.py
@@ -14,7 +14,7 @@
     pipe.model,
     mp_size=world_size,
     dtype=torch.float,
-    injection_policy={BertLayer : ('output.dense')}
+    replace_with_kernel_inject=True
 )
 
 pipe.device = torch.device(f'cuda:{local_rank}')

Original file line number	Diff line number	Diff line change
`@@ -14,7 +14,7 @@`
`14`	`14`	`pipe.model,`
`15`	`15`	`mp_size=world_size,`
`16`	`16`	`dtype=torch.float,`
`17`		`- injection_policy={BertLayer : ('output.dense')}`
	`17`	`+ replace_with_kernel_inject=True`
`18`	`18`	`)`
`19`	`19`
`20`	`20`	`pipe.device = torch.device(f'cuda:{local_rank}')`