CPU and M1/M2 GPU platform support

README.md

@@ -131,6 +131,13 @@ You can then try to offloading all weights to disk by
 python3 -m flexgen.flex_opt --model facebook/opt-175b --percent 0 0 100 0 100 0 --offload-dir YOUR_SSD_FOLDER
 ```
 
+### CPU and M1/M2 GPU platform
+To run models on CPU platforms, all you need to do is to add an `--platform` entry:
+```
+python3 -m flexgen.flex_opt --model facebook/opt-1.3b --platform cpu
+```
+To run on M1/M2 platforms, [PyTorch nightly](https://pytorch.org/) is required for kernel coverage and better performance. Once you have PyTorch nightly installed, you can simply replace `cpu` with `mps:0`.
+
 ### How to set the offloading strategy and `--percent`?
 We will release an automatic policy optimizer later, but now you have to manually try a few strategies.
 The idea of high-throughput generation is to offload parameters and attention cache as much as possible to the CPU and disk if necessary.
@@ -191,7 +198,7 @@ See [flexgen/apps](flexgen/apps) for more example applications.
 ## Roadmap
 We plan to work on the following features. Community contributions are welcome.
 
-- [ ] Support Apple silicon M1/M2 deployment
+- [x] Support Apple silicon M1/M2 deployment
 - [ ] Support Colab deployment
 - [ ] Support more models (BLOOM, CodeGen, GLM)
 - [ ] Release the cost model and policy optimizer

flexgen/flex_opt.py

@@ -43,6 +43,8 @@ class Policy:
     act_gpu_percent: float
     act_cpu_percent: float
 
+    only_cpu: bool
+
     # Whether to overlap the I/O and compute
     overlap: bool
 
@@ -101,12 +103,8 @@ def init_weight_list(weight_specs, policy, env):
         home = get_choice(mid_percent * 100, dev_percents, dev_choices)
         shape, dtype, filename = weight_specs[i]
 
-        if len(shape) < 2:
-            pin_memory = True
-            compress = False
-        else:
-            pin_memory = policy.pin_weight
-            compress = policy.compress_weight
+        pin_memory = policy.pin_weight
+        compress = policy.compress_weight
 
         if not compress:
             weight = home.allocate(shape, dtype, pin_memory=pin_memory)
@@ -614,10 +612,11 @@ def __init__(self,
         else:
             raise NotImplementedError()
 
-        # CUDA streams
-        self.load_weight_stream = torch.cuda.Stream()
-        self.load_cache_stream = torch.cuda.Stream()
-        self.store_cache_stream = torch.cuda.Stream()
+        if self.env.gpu.device_type == DeviceType.CUDA:
+            # CUDA streams
+            self.load_weight_stream = torch.cuda.Stream()
+            self.load_cache_stream = torch.cuda.Stream()
+            self.store_cache_stream = torch.cuda.Stream()
 
         # Intermediate tensors
         # The following buffers store values used
@@ -791,7 +790,8 @@ def compute_layer(self, i, j, k):
 
     def sync(self):
         self.env.disk.synchronize()
-        torch.cuda.synchronize()
+        if self.env.gpu.device_type == DeviceType.CUDA:
+            torch.cuda.synchronize()
 
     def init_all_weights(self):
         self.weight_home = array_1d(self.num_layers, ValueHolder)
@@ -1184,15 +1184,17 @@ def run_flexgen(args):
     warmup_inputs = get_test_inputs(32, num_prompts, tokenizer)
     inputs = get_test_inputs(prompt_len, num_prompts, tokenizer)
 
-    gpu = TorchDevice("cuda:0")
+    if args.platform == "cpu":
+        gpu = TorchDevice("cpu")
+    else:
+        gpu = TorchDevice(args.platform)
     cpu = TorchDevice("cpu")
-    disk = TorchDisk(args.offload_dir)
-    env = ExecutionEnv(gpu=gpu, cpu=cpu, disk=disk, mixed=TorchMixedDevice([gpu, cpu, disk]))
-
+    disk = TorchDisk(args.offload_dir,platform=args.platform)
+    env = ExecutionEnv(gpu=gpu, cpu=cpu, disk=disk, mixed=TorchMixedDevice([gpu, cpu, disk]), platform=args.platform)
     policy = Policy(args.gpu_batch_size, args.num_gpu_batches,
                     args.percent[0], args.percent[1],
                     args.percent[2], args.percent[3],
-                    args.percent[4], args.percent[5],
+                    args.percent[4], args.percent[5], args.platform == "cpu",
                     args.overlap, args.sep_layer, args.pin_weight,
                     args.cpu_cache_compute, args.attn_sparsity,
                     args.compress_weight,
@@ -1203,7 +1205,8 @@ def run_flexgen(args):
                                       group_dim=2, symmetric=False))
     assert not (args.compress_cache and args.attn_sparsity < 1.0), "Not implemented"
 
-    opt_config = get_opt_config(args.model)
+    # use float32 for CPU platform
+    opt_config = get_opt_config(args.model, dtype=np.float32 if args.platform == "cpu" else np.float16)
     cache_size = opt_config.cache_bytes(num_prompts, prompt_len + gen_len)
     hidden_size = opt_config.hidden_bytes(num_prompts, prompt_len + gen_len)
     print(f"model size: {opt_config.model_bytes()/GB:.3f} GB, "
@@ -1311,6 +1314,7 @@ def add_parser_arguments(parser):
     parser.add_argument("--overlap", type=str2bool, nargs='?',
         const=True, default=True)
 
+    parser.add_argument("--platform", type=str, default="cuda:0", help="use the number to specify device, the platform can also be cpu or mps")
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -1319,4 +1323,25 @@ def add_parser_arguments(parser):
 
     assert len(args.percent) == 6
 
+    if "cuda" in args.platform:
+        if not torch.cuda.is_available():
+            if torch.backends.mps.is_available():
+                args.platform = "mps:0"
+            else:
+                args.platform = "cpu"
+            print("CUDA devices not available, {} is used instead".format(args.platform))
+
+    if "mps" in args.platform:
+        if not torch.backends.mps.is_available():
+            args.platform = "cpu"
+            print("MPS devices not available, CPU is used instead")
+
+    if "cuda" not in args.platform:
+        # not clear how to enable overlap on MPS platform yet
+        args.overlap = False
+        args.pin_weight = False
+
+    if args.platform == "cpu":
+        args.percent = [0, 100, 0, 100, 0, 100]
+
     run_flexgen(args)

flexgen/opt_config.py

@@ -54,6 +54,8 @@ def get_opt_config(name, **kwargs):
         name = name.split("/")[1]
     name = name.lower()
 
+    dtype = kwargs["dtype"] if "dtype" in kwargs else OptConfig.dtype
+
     # Handle opt-iml-30b and opt-iml-max-30b
     if "-iml-max" in name:
         arch_name = name.replace("-iml-max", "")
@@ -65,54 +67,54 @@ def get_opt_config(name, **kwargs):
     if arch_name == "opt-125m":
         config = OptConfig(name=name,
             max_seq_len=2048, num_hidden_layers=12, n_head=12,
-            hidden_size=768, input_dim=768, ffn_embed_dim=768 * 4,
+            hidden_size=768, input_dim=768, ffn_embed_dim=768 * 4, dtype=dtype
         )
     elif arch_name == "opt-350m":
         config = OptConfig(name=name,
             max_seq_len=2048, num_hidden_layers=24, n_head=16,
-            hidden_size=1024, input_dim=1024, ffn_embed_dim=1024 * 4,
+            hidden_size=1024, input_dim=1024, ffn_embed_dim=1024 * 4, dtype=dtype
         )
         raise NotImplementedError("Not implemented because this model "
                                   "has a different architecture")
     elif arch_name == "opt-1.3b":
         config = OptConfig(name=name,
             max_seq_len=2048, num_hidden_layers=24, n_head=32,
-            hidden_size=2048, input_dim=2048, ffn_embed_dim=2048 * 4,
+            hidden_size=2048, input_dim=2048, ffn_embed_dim=2048 * 4, dtype=dtype
         )
     elif arch_name == "opt-2.7b":
         config = OptConfig(name=name,
             max_seq_len=2048, num_hidden_layers=32, n_head=32,
-            hidden_size=2560, input_dim=2560, ffn_embed_dim=2560 * 4,
+            hidden_size=2560, input_dim=2560, ffn_embed_dim=2560 * 4, dtype=dtype
         )
     elif arch_name == "opt-6.7b":
         config = OptConfig(name=name,
             max_seq_len=2048, num_hidden_layers=32, n_head=32,
-            hidden_size=4096, input_dim=4096, ffn_embed_dim=4096 * 4,
+            hidden_size=4096, input_dim=4096, ffn_embed_dim=4096 * 4, dtype=dtype
         )
     elif arch_name == "opt-13b":
         config = OptConfig(name=name,
             max_seq_len=2048, num_hidden_layers=40, n_head=40,
-            hidden_size=5120, input_dim=5120, ffn_embed_dim=5120 * 4,
+            hidden_size=5120, input_dim=5120, ffn_embed_dim=5120 * 4, dtype=dtype
         )
     elif arch_name == "opt-30b":
         config = OptConfig(name=name,
             max_seq_len=2048, num_hidden_layers=48, n_head=56,
-            hidden_size=7168, input_dim=7168, ffn_embed_dim=7168 * 4,
+            hidden_size=7168, input_dim=7168, ffn_embed_dim=7168 * 4, dtype=dtype
         )
     elif arch_name == "opt-66b":
         config = OptConfig(name=name,
             max_seq_len=2048, num_hidden_layers=64, n_head=72,
-            hidden_size=9216, input_dim=9216, ffn_embed_dim=9216 * 4,
+            hidden_size=9216, input_dim=9216, ffn_embed_dim=9216 * 4, dtype=dtype
         )
     elif arch_name == "opt-175b":
         config = OptConfig(name=name,
             max_seq_len=2048, num_hidden_layers=96, n_head=96,
-            hidden_size=12288, input_dim=12288, ffn_embed_dim=12288 * 4,
+            hidden_size=12288, input_dim=12288, ffn_embed_dim=12288 * 4, dtype=dtype
         )
     elif arch_name == "opt-175b-stage":
         config = OptConfig(name=name,
             max_seq_len=2048, num_hidden_layers=24, n_head=96,
-            hidden_size=12288, input_dim=12288, ffn_embed_dim=12288 * 4,
+            hidden_size=12288, input_dim=12288, ffn_embed_dim=12288 * 4, dtype=dtype
         )
     else:
         raise ValueError(f"Invalid model name: {name}")