Jonathan1909 commited on
Commit
5a913a1
1 Parent(s): 7546e4f

Add README

Browse files
Files changed (1) hide show
  1. README.md +47 -0
README.md ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ pipeline_tag: text-generation
4
+ ---
5
+ # Grok-1 (PyTorch Version)
6
+
7
+ This repository contains the model and weights of the **torch version** of Grok-1 open-weights model.
8
+
9
+ You could find the original weights released by [xAI](https://x.ai/blog) in [Hugging Face](https://huggingface.co/xai-org/grok-1) and the original model in the Grok open release [GitHub Repository](https://github.com/xai-org/grok-1/tree/main).
10
+
11
+ ## Conversion
12
+
13
+ We translated the original modeling written in JAX into PyTorch version, and converted the weights by mapping tensor files with parameter keys, de-quantizing the tensors with corresponding packed scales, and save to checkpoint file with torch APIs.
14
+
15
+ The original tokenizer is supposed to be used (i.e. `tokenizer.model` in [GitHub Repository](https://github.com/xai-org/grok-1/tree/main)) with the torch-version model.
16
+
17
+ ## Usage
18
+
19
+ ```python
20
+ import torch
21
+ from transformers import AutoModelForCausalLM
22
+
23
+ torch.set_default_dtype(torch.bfloat16)
24
+ model = AutoModelForCausalLM.from_pretrained(
25
+ "hpcaitech/grok-1",
26
+ trust_remote_code=True,
27
+ device_map="auto",
28
+ torch_dtype=torch.bfloat16,
29
+ )
30
+ sp = SentencePieceProcessor(model_file="tokenizer.model")
31
+
32
+ text = "Replace this with your text"
33
+ input_ids = sp.encode(text)
34
+ input_ids = torch.tensor([input_ids]).cuda()
35
+ attention_mask = torch.ones_like(input_ids)
36
+ generate_kwargs = {} # Add any additional args if you want
37
+ inputs = {
38
+ "input_ids": input_ids,
39
+ "attention_mask": attention_mask,
40
+ **generate_kwargs,
41
+ }
42
+ outputs = model.generate(**inputs)
43
+ ```
44
+
45
+ You could find a complete example code of using the torch-version Grok-1 in ColossalAI [GitHub Repository](https://github.com/hpcaitech/ColossalAI/tree/main/examples/language/grok-1). We also applies parallelism techniques from ColossalAI framework (Tensor Parallelism for now) to accelerate the inference.
46
+
47
+ Note: A multi-GPU machine is required to test the model with the example code (For now, a 8x80G multi-GPU machine is required).