├── .DS_Store
├── DATASET.md
├── EasyLM
├── __init__.py
├── bpt.py
├── checkpoint.py
├── data.py
├── jax_utils.py
├── models
│ ├── __init__.py
│ └── llama
│ │ ├── convert_easylm_to_hf.py
│ │ ├── llama_model.py
│ │ ├── llama_serve.py
│ │ └── llama_train.py
├── optimizers.py
├── scripts
│ ├── __init__.py
│ ├── convert_checkpoint.py
│ ├── diff_checkpoint.py
│ ├── lm_eval_harness.py
│ └── lm_eval_json.py
└── serving.py
├── LICENSE
├── README.md
├── evaluation
├── .DS_Store
├── EVAL.md
├── LICENSE
├── environment.yml
└── vqlm_demo
│ ├── .DS_Store
│ ├── __init__.py
│ ├── batch_generation.py
│ ├── eval_perplexity.py
│ ├── eval_video_perplexity.py
│ ├── eval_videos.py
│ ├── generate_videos.py
│ ├── inference.py
│ ├── torch_vqvae_model.py
│ ├── utils.py
│ ├── vqvae
│ ├── .DS_Store
│ ├── __init__.py
│ ├── logging.py
│ └── modeling_utils.py
│ └── vqvae_muse.py
├── images
└── visual_sentences.jpg
├── scripts
├── gpu_environment.yml
├── tpu_commands.sh
└── tpu_vm_setup.sh
└── tokenize_examples
├── detokenization_muse.py
├── map_color.py
├── tokenize_catogory_images_muse.py
├── tokenize_co3d_muse.py
├── tokenize_colorization_dataset_muse.py
├── tokenize_inpainting_dataset_muse.py
├── tokenize_multi_datasets_muse.py
├── tokenize_multi_seq_images_muse.py
├── tokenize_paired_dataset_muse.py
├── tokenize_seq_images_muse.py
└── tokenize_video_muse.py
/.DS_Store:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/ytongbai/LVM/b6de939ef0eb1ee6593445a7f5268145f338749b/.DS_Store
--------------------------------------------------------------------------------
/DATASET.md:
--------------------------------------------------------------------------------
1 |
2 | # Dataset Preparation
3 |
4 | This section describes how to prepare your dataset by tokenizing visual data into visual sentences and then mixing and shuffling the datasets.
5 |
6 | ## Download & Prepare Dataset
7 | We have individual scripts in `./tokenize_examples` to handle different kinds of visual sentences mentioned in the paper. Simple descriptions and instructions are listed as follows.
8 |
9 | ### Pair Datasets
10 | For pair datasets, where the visual sentence is constructed as `[image, label, image, label, image, label...]`, we use `tokenize_examples/tokenize_paired_dataset_muse.py` to generate the visual tokens. For example, for depth maps, surface normals, edges, and segmentation, we first use [prismer](https://github.com/NVlabs/prismer) to generate pseudo labels, then use the script to generate visual sentences. Note that for segmentation, we use an additional color mapping after obtaining the pseudo labels from prismer, which can be done by running `tokenize_examples/map_color.py`.
11 |
12 | ### Video Datasets
13 | For video datasets, the visual sentences are constructed as `[frame1, frame2, frame3, ... framex]`. One can use `tokenize_examples/tokenize_video_muse.py` to generate the visual sentences. The hyperparameter `stride` can be used to control the sampling rate of the extraction of frames from a video.
14 |
15 | ### Colorization Datasets
16 | For colorization datasets, the visual sentences are constructed as `[gray_image, colored_image, gray_image, colored_image, ...]`. One can use `tokenize_examples/tokenize_colorization_dataset_muse.py` to generate the visual sentences. The user only needs to prepare the colored images, and the script will generate the gray counterparts.
17 |
18 | ### Inpainting Datasets
19 | For inpainting datasets, the visual sentences are constructed as `[masked_image, image, masked_image, image, ...]`. One can use `tokenize_examples/tokenize_inpainting_dataset_muse.py` to generate the visual sentences. The user can control the masked ratio by changing `FLAGS.hole_mask_ratio`.
20 |
21 | ### Multi-Datasets
22 | For multi-datasets, the visual sentences are constructed as `[dataset1_image, dataset1_image, dataset2_image, dataset2_image, ...]`. One can use `tokenize_examples/tokenize_multi_datasets_muse.py` to generate the visual sentences.
23 |
24 | ### Category Datasets
25 | For category datasets, the visual sentences are constructed as `[cate1_image, cate1_image, cate2_image, cate2_image,...]`. One can use `tokenize_examples/tokenize_category_images_muse.py` to generate the visual sentences. Note that the user can use `images_per_shot` to customize the number of images for each category and `n_shots` to control the number of categories in the visual sentences.
26 |
27 | In general, visual sentences with different logic can be achieved by combining the basic logic provided above. After you generate your visual sentences, you can always use `tokenize_examples/detokenization_muse.py` to perform a sanity check for the recovery of the visual sentences.
28 |
29 |
30 | ## Visual Sentences
31 |
32 | - Following the concept of Visual Sentences, each image is tokenized separately, and the tokens are concatenated to form a visual sentence.
33 |
34 | ## Generating JSONL Files
35 |
36 | - For each dataset, generate the `dataset*.jsonl` files.
37 |
38 | ### Example Code
39 |
40 | - You can find example code for tokenizing images in the `./tokenize_examples` directory.
41 |
42 | ## Mixing and Shuffling Datasets
43 |
44 | - After generating the JSONL files, the datasets need to be mixed and shuffled. Follow these steps:
45 |
46 | 1. Set the temporary directory and memory allocation:
47 | ```shell
48 | export TMPDIR='/global/scratch/users/yutong/data/temp'
49 | export MEMORY='20'
50 | ```
51 |
52 | 2. Navigate to your data directory:
53 | ```shell
54 | cd /path/to/your/data/
55 | ```
56 |
57 | 3. Mix and shuffle the datasets:
58 | ```shell
59 | cat tokenized_tasks/*.jsonl | terashuf > mix_and_shuffled/dataset.jsonl
60 | ```
61 |
62 | - This will create a mixed and shuffled dataset file named `dataset.jsonl` in the `mix_and_shuffled` directory.
--------------------------------------------------------------------------------
/EasyLM/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/ytongbai/LVM/b6de939ef0eb1ee6593445a7f5268145f338749b/EasyLM/__init__.py
--------------------------------------------------------------------------------
/EasyLM/bpt.py:
--------------------------------------------------------------------------------
1 | """An implementation of Blockwise parallel transformer https://arxiv.org/abs/2305.19370
2 |
3 | Also include a reference implementation of memory-efficient transformer https://arxiv.org/abs/2112.05682
4 | """
5 |
6 | import functools
7 | from typing import NamedTuple
8 |
9 | import flax.linen as nn
10 | import jax
11 | import jax.lax as lax
12 | import jax.numpy as jnp
13 | from einops import rearrange
14 |
15 | '''
16 | Computing ffn blockwise without materializing the large hidden tensor, training 4x longer sequences than the memory-efficient transformer.
17 | Blockwise parallel transformer https://arxiv.org/abs/2305.19370 Liu et al. 2023
18 | '''
19 | def blockwise_ffn(remat_ffn, inputs, chunk_size, deterministic):
20 | # remat_ffn: a rematerialized ffn with policy jax.checkpoint_policies.nothing_saveable()
21 | # inputs: (batch, seq_len, dim)
22 | # chunk_size: the chunk size to split the sequence
23 | inputs = rearrange(inputs, 'b (c n) d -> b c n d', c=chunk_size)
24 | def scan_ffn(remat_ffn, carry, hidden_states):
25 | outputs = remat_ffn(hidden_states, deterministic=deterministic)
26 | return carry, outputs
27 | scan_axis = inputs.ndim - 2
28 | _, res = nn.scan(
29 | scan_ffn,
30 | variable_broadcast="params",
31 | split_rngs={"params": False, "dropout": True},
32 | in_axes=scan_axis,
33 | out_axes=scan_axis,
34 | )(remat_ffn, None, inputs)
35 | res = rearrange(res, 'b c n d -> b (c n) d')
36 | return res
37 |
38 |
39 | '''
40 | Compute attention blockwise without materializing the full attention matrix, initially proposed in memory-efficient transformer https://arxiv.org/abs/2112.05682 Rabe et al. 2021;
41 | flash attention https://arxiv.org/abs/2205.14135 Dao et al. 2022 proposes a CUDA efficient implementation;
42 | blockwise parallel transformer https://arxiv.org/abs/2305.19370 Liu et al. 2023 proposes blockwise computing both attention and FFN, enabling 4x longer sequences than memory-efficient/flash-attention and fusion of attention and FFN.
43 | '''
44 | def blockwise_attn(query, key, value, bias, deterministic,
45 | dropout_rng, attn_pdrop, causal, query_chunk_size,
46 | key_chunk_size, dtype, policy, precision, float32_logits,
47 | prevent_cse):
48 | # query, key, value: (batch, seq_len, num_heads, dim_per_head)
49 | # bias: (batch, seq_len) can be used to mask out attention (e.g. padding)
50 | # causal: whether to use causal mask
51 | # policy: one of jax.checkpoint_policies
52 | query = query / jnp.sqrt(query.shape[-1]).astype(dtype)
53 | if float32_logits:
54 | query = query.astype(jnp.float32)
55 | key = key.astype(jnp.float32)
56 |
57 | batch, q_len, num_heads, dim_per_head = query.shape
58 | batch, kv_len, num_heads, dim_per_head = key.shape
59 | batch, kv_len, num_heads, dim_per_head = value.shape
60 |
61 | num_q = q_len // query_chunk_size
62 | num_kv = kv_len // key_chunk_size
63 | query = query.reshape((batch, num_q, query_chunk_size, num_heads, dim_per_head))
64 | key = key.reshape((batch, num_kv, key_chunk_size, num_heads, dim_per_head))
65 | value = value.reshape((batch, num_kv, key_chunk_size, num_heads, dim_per_head))
66 |
67 | query = jnp.moveaxis(query, 1, 0)
68 | key = jnp.moveaxis(key, 1, 0)
69 | value = jnp.moveaxis(value, 1, 0)
70 |
71 | if bias is not None:
72 | for bias_dim, broadcast_dim in zip(bias.shape, (batch, num_heads, q_len, kv_len)):
73 | assert bias_dim == 1 or bias_dim == broadcast_dim
74 | if not deterministic and attn_pdrop > 0.0:
75 | attn_dropout_rng, dropout_rng = jax.random.split(dropout_rng)
76 | attn_dropout = jax.random.bernoulli(attn_dropout_rng, attn_pdrop, (batch, num_heads, q_len, kv_len))
77 | else:
78 | attn_dropout = None
79 |
80 | _chunk_bias_fn = functools.partial(
81 | _chunk_attention_bias,
82 | query_chunk_size, key_chunk_size, bias, deterministic,
83 | attn_dropout, attn_pdrop, causal, dtype)
84 |
85 | def scan_attention(args):
86 | query_chunk, query_chunk_idx = args
87 |
88 | @functools.partial(jax.checkpoint, prevent_cse=prevent_cse, policy=policy)
89 | def scan_kv_block(carry, args):
90 | key_chunk, value_chunk, key_chunk_idx = args
91 | (numerator, denominator, prev_max_score) = carry
92 | attn_weights = jnp.einsum('bqhd,bkhd->bqhk', query_chunk, key_chunk, precision=precision)
93 | bias_chunk = _chunk_bias_fn(query_chunk_idx, key_chunk_idx)
94 | bias_chunk = jnp.moveaxis(bias_chunk, 1, 2)
95 | attn_weights = attn_weights + bias_chunk
96 |
97 | max_score = jnp.max(attn_weights, axis=-1, keepdims=True)
98 | max_score = jnp.maximum(prev_max_score, max_score)
99 | max_score = jax.lax.stop_gradient(max_score)
100 | exp_weights = jnp.exp(attn_weights - max_score)
101 | exp_values = jnp.einsum(
102 | 'bqhv,bvhd->bqhd', exp_weights, value_chunk, precision=precision
103 | )
104 | correction = jnp.exp(prev_max_score - max_score)
105 | numerator = numerator * correction + exp_values
106 | denominator = denominator * correction + exp_weights.sum(axis=-1, keepdims=True)
107 | return Carry(numerator, denominator, max_score), None
108 |
109 | def skip_upper_half(carry, args):
110 | key_chunk, value_chunk, key_chunk_idx = args
111 | skip_block = jnp.array(False)
112 | if causal:
113 | skip_block = query_chunk_idx < key_chunk_idx
114 | return jax.lax.cond(
115 | skip_block,
116 | lambda carry, args: (carry, None),
117 | scan_kv_block,
118 | carry,
119 | args,
120 | )
121 |
122 | init_carry = Carry(
123 | jnp.zeros((batch, query_chunk_size, num_heads, dim_per_head), dtype=query.dtype),
124 | jnp.zeros((batch, query_chunk_size, num_heads, dim_per_head), dtype=query.dtype),
125 | (-jnp.inf) * jnp.ones((batch, query_chunk_size, num_heads, 1), dtype=query.dtype),
126 | )
127 | (numerator, denominator, max_score), _ = lax.scan(
128 | skip_upper_half, init_carry, xs=(key, value, jnp.arange(0, num_kv))
129 | )
130 | outputs = (numerator / denominator).astype(dtype)
131 | return outputs
132 |
133 | _, res = lax.scan(
134 | lambda _, x: ((), scan_attention(x)),
135 | (), xs=(query, jnp.arange(0, num_q))
136 | )
137 | res = rearrange(res, 'n b c h d -> b (n c) h d')
138 | return res
139 |
140 |
141 | class Carry(NamedTuple):
142 | numerator: jax.Array
143 | denominator: jax.Array
144 | max_so_far: jax.Array
145 |
146 |
147 | def _chunk_attention_bias(query_chunk_size, key_chunk_size,
148 | bias, deterministic, attn_dropout, attn_pdrop, causal,
149 | dtype, query_chunk_idx, key_chunk_idx):
150 | query_offset = query_chunk_idx * query_chunk_size
151 | key_offset = key_chunk_idx * key_chunk_size
152 | chunk_bias = jnp.zeros((1, 1, 1, 1), dtype=dtype)
153 | if bias is not None:
154 | chunk_bias = lax.dynamic_slice(
155 | bias,
156 | start_indices=(0, 0, query_offset, key_offset),
157 | slice_sizes=(*bias.shape[:2], min(bias.shape[-2], query_chunk_size), min(bias.shape[-1], key_chunk_size)),
158 | )
159 |
160 | if causal:
161 | query_idx = lax.broadcasted_iota(dtype=jnp.int32, shape=(query_chunk_size, 1), dimension=0)
162 | key_idx = lax.broadcasted_iota(dtype=jnp.int32, shape=(1, key_chunk_size), dimension=1)
163 | offset = query_offset - key_offset
164 | query_idx += offset
165 | causal_mask_value = (query_idx < key_idx) * jnp.finfo(dtype).min
166 | chunk_bias += causal_mask_value.reshape(1, 1, *causal_mask_value.shape)
167 |
168 | if not deterministic and attn_pdrop > 0.0:
169 | attn_dropout_slice = lax.dynamic_slice(
170 | attn_dropout,
171 | start_indices=(0, 0, query_offset, key_offset),
172 | slice_sizes=(
173 | *attn_dropout.shape[:2],
174 | min(attn_dropout.shape[-2], query_chunk_size),
175 | min(attn_dropout.shape[-1], key_chunk_size),
176 | ),
177 | )
178 | chunk_bias += attn_dropout_slice * jnp.finfo(dtype).min
179 | return chunk_bias.astype(dtype)
180 |
181 |
182 | if __name__ == '__main__':
183 | # test
184 | def reference_attn(query, key, value, causal, dtype):
185 | query = query / jnp.sqrt(query.shape[-1]).astype(dtype)
186 | logits = jnp.einsum("bqhc,bkhc->bhqk", query, key)
187 | if causal:
188 | mask_value = jnp.finfo(logits.dtype).min
189 | _, q_seq_len, _, _ = query.shape
190 | _, kv_seq_len, _, _ = key.shape
191 | mask_shape = (q_seq_len, kv_seq_len)
192 | row_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 0)
193 | col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1)
194 | causal_mask = (row_ids < col_ids)[None, None, :, :]
195 | logits = logits + jnp.where(causal_mask, mask_value, 0.0)
196 | weights = jax.nn.softmax(logits, axis=-1)
197 | out = jnp.einsum("bhqk,bkhc->bqhc", weights, value)
198 | return out
199 |
200 | # random inputs
201 | shape = (1, 32, 8, 64)
202 | query = jax.random.normal(jax.random.PRNGKey(0), shape)
203 | key = jax.random.normal(jax.random.PRNGKey(1), shape)
204 | value = jax.random.normal(jax.random.PRNGKey(2), shape)
205 |
206 | causal = True
207 | chunk_size = 4
208 | policy = jax.checkpoint_policies.nothing_saveable()
209 |
210 | blockwise = blockwise_attn(query, key, value, None, False, None, 0.0, causal, chunk_size, chunk_size, jnp.float32, policy, 'float32', True, False)
211 | reference = reference_attn(query, key, value, causal, 'float32')
212 |
213 | assert jnp.allclose(reference, blockwise, atol=1e-6)
214 |
--------------------------------------------------------------------------------
/EasyLM/checkpoint.py:
--------------------------------------------------------------------------------
1 | import os
2 | import numpy as np
3 | from ml_collections import ConfigDict
4 | import mlxu
5 | import jax
6 | import jax.numpy as jnp
7 | import flax
8 | from flax.serialization import (
9 | from_bytes, to_bytes, to_state_dict, from_state_dict
10 | )
11 | from flax.traverse_util import flatten_dict, unflatten_dict, empty_node
12 | import msgpack
13 |
14 | from EasyLM.jax_utils import tree_apply, float_tensor_to_dtype
15 |
16 |
17 | class StreamingCheckpointer(object):
18 | """ Custom msgpack checkpointer that saves large train states by serializing
19 | and saving tensors one by one in a streaming fashion. Avoids running
20 | out of memory or local TPU disk with default flax checkpointer.
21 | """
22 |
23 | @staticmethod
24 | def get_default_config(updates=None):
25 | config = ConfigDict()
26 | config.float_dtype = 'bf16'
27 | config.save_optimizer_state = False
28 |
29 | if updates is not None:
30 | config.update(ConfigDict(updates).copy_and_resolve_references())
31 | return config
32 |
33 | def __init__(self, config, checkpoint_dir, enable=True):
34 | self.config = self.get_default_config(config)
35 | self.checkpoint_dir = checkpoint_dir
36 | self.enable = enable
37 |
38 | def save_checkpoint(self, train_state, filename, gather_fns=None):
39 | if self.enable:
40 | path = os.path.join(self.checkpoint_dir, filename)
41 | else:
42 | path = '/dev/null'
43 | self.save_train_state_to_file(
44 | train_state, path, gather_fns, self.config.float_dtype
45 | )
46 |
47 | @staticmethod
48 | def save_train_state_to_file(train_state, path, gather_fns=None, float_dtype=None):
49 | train_state = to_state_dict(train_state)
50 | packer = msgpack.Packer()
51 | flattend_train_state = flatten_dict(train_state)
52 | if gather_fns is not None:
53 | gather_fns = flatten_dict(to_state_dict(gather_fns))
54 |
55 | with mlxu.open_file(path, "wb") as fout:
56 | for key, value in flattend_train_state.items():
57 | if gather_fns is not None:
58 | value = gather_fns[key](value)
59 | value = float_tensor_to_dtype(value, float_dtype)
60 | fout.write(packer.pack((key, to_bytes(value))))
61 |
62 | def save_pickle(self, obj, filename):
63 | if self.enable:
64 | path = os.path.join(self.checkpoint_dir, filename)
65 | else:
66 | path = '/dev/null'
67 | mlxu.save_pickle(obj, path)
68 |
69 | def save_all(self, train_state, gather_fns, metadata=None, dataset=None, milestone=False):
70 | step = int(jax.device_get(train_state.step))
71 | if self.config.save_optimizer_state:
72 | checkpoint_state = train_state
73 | checkpoint_name = 'streaming_train_state'
74 | checkpoint_gather_fns = gather_fns
75 | else:
76 | checkpoint_state = train_state.params['params']
77 | checkpoint_name = 'streaming_params'
78 | checkpoint_gather_fns = gather_fns.params['params']
79 |
80 | if milestone:
81 | # Save a milestone checkpoint that will not be overwritten
82 | self.save_pickle(metadata, f'metadata_{step}.pkl')
83 | self.save_pickle(dataset, f'dataset_{step}.pkl')
84 | self.save_checkpoint(
85 | checkpoint_state, f'{checkpoint_name}_{step}', checkpoint_gather_fns
86 | )
87 | else:
88 | # Save a normal checkpoint that can be overwritten
89 | self.save_pickle(metadata, 'metadata.pkl')
90 | self.save_pickle(dataset, 'dataset.pkl')
91 | self.save_checkpoint(
92 | checkpoint_state, f'{checkpoint_name}', checkpoint_gather_fns
93 | )
94 |
95 | @staticmethod
96 | def load_checkpoint(path, target=None, shard_fns=None, remove_dict_prefix=None):
97 | if shard_fns is not None:
98 | shard_fns = flatten_dict(
99 | to_state_dict(shard_fns)
100 | )
101 | if remove_dict_prefix is not None:
102 | remove_dict_prefix = tuple(remove_dict_prefix)
103 | flattend_train_state = {}
104 | with mlxu.open_file(path) as fin:
105 | # 83886080 bytes = 80 MB, which is 16 blocks on GCS
106 | unpacker = msgpack.Unpacker(fin, read_size=83886080, max_buffer_size=0)
107 | for key, value in unpacker:
108 | key = tuple(key)
109 | if remove_dict_prefix is not None:
110 | if key[:len(remove_dict_prefix)] == remove_dict_prefix:
111 | key = key[len(remove_dict_prefix):]
112 | else:
113 | continue
114 |
115 | tensor = from_bytes(None, value)
116 | if shard_fns is not None:
117 | tensor = shard_fns[key](tensor)
118 | flattend_train_state[key] = tensor
119 |
120 | if target is not None:
121 | flattened_target = flatten_dict(
122 | to_state_dict(target), keep_empty_nodes=True
123 | )
124 | for key, value in flattened_target.items():
125 | if key not in flattend_train_state and value == empty_node:
126 | flattend_train_state[key] = value
127 |
128 | train_state = unflatten_dict(flattend_train_state)
129 | if target is None:
130 | return train_state
131 |
132 | return from_state_dict(target, train_state)
133 |
134 | @staticmethod
135 | def load_flax_checkpoint(path, target=None, shard_fns=None):
136 | """ Load a standard flax checkpoint that's not saved with the
137 | msgpack streaming format.
138 | """
139 | with mlxu.open_file(path, "rb") as fin:
140 | encoded_bytes = fin.read()
141 |
142 | state_dict = flax.serialization.msgpack_restore(encoded_bytes)
143 | if shard_fns is not None:
144 | shard_fns = to_state_dict(shard_fns)
145 | state_dict = tree_apply(shard_fns, state_dict)
146 |
147 | if target is None:
148 | return state_dict
149 | return from_state_dict(target, state_dict)
150 |
151 | @classmethod
152 | def load_trainstate_checkpoint(cls, load_from, trainstate_target=None,
153 | trainstate_shard_fns=None,
154 | disallow_trainstate=False):
155 | if trainstate_target is not None:
156 | params_target = trainstate_target.params['params']
157 | else:
158 | params_target = None
159 |
160 | if trainstate_shard_fns is not None:
161 | params_shard_fns = trainstate_shard_fns.params['params']
162 | else:
163 | params_shard_fns = None
164 |
165 | load_type, load_path = load_from.split('::', 1)
166 | if disallow_trainstate:
167 | assert load_type != 'trainstate', 'Loading full trainstate is not allowed!'
168 | train_state = None
169 | restored_params = None
170 | if load_type == 'trainstate':
171 | # Load the entire train state in the streaming format
172 | train_state = cls.load_checkpoint(
173 | path=load_path,
174 | target=trainstate_target,
175 | shard_fns=trainstate_shard_fns,
176 | )
177 | elif load_type == 'trainstate_params':
178 | # Load the params part of the train state in the streaming format
179 | restored_params = cls.load_checkpoint(
180 | path=load_path,
181 | target=params_target,
182 | shard_fns=params_shard_fns,
183 | remove_dict_prefix=('params', 'params'),
184 | )
185 | restored_params = flax.core.frozen_dict.freeze(
186 | {'params': restored_params}
187 | )
188 | elif load_type == 'params':
189 | # Load the params in the streaming format
190 | restored_params = cls.load_checkpoint(
191 | path=load_path,
192 | target=params_target,
193 | shard_fns=params_shard_fns,
194 | )
195 | restored_params = flax.core.frozen_dict.freeze(
196 | {'params': restored_params}
197 | )
198 | elif load_type == 'flax_params':
199 | # Load the params in the standard flax format (non-streaming)
200 | # This requires the entire params to fit in memory
201 | restored_params = cls.load_flax_checkpoint(
202 | path=load_path,
203 | target=params_target,
204 | shard_fns=params_shard_fns
205 | )
206 | restored_params = flax.core.frozen_dict.freeze(
207 | {'params': restored_params}
208 | )
209 | else:
210 | raise ValueError(f'Invalid load_from type: {load_type}')
211 |
212 | return train_state, restored_params
213 |
--------------------------------------------------------------------------------
/EasyLM/models/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/ytongbai/LVM/b6de939ef0eb1ee6593445a7f5268145f338749b/EasyLM/models/__init__.py
--------------------------------------------------------------------------------
/EasyLM/models/llama/convert_easylm_to_hf.py:
--------------------------------------------------------------------------------
1 | # Copyright 2022 EleutherAI and The HuggingFace Inc. team. All rights reserved.
2 | # Copyright 2023 Xinyang Geng
3 | #
4 | # Licensed under the Apache License, Version 2.0 (the "License");
5 | # you may not use this file except in compliance with the License.
6 | # You may obtain a copy of the License at
7 | #
8 | # http://www.apache.org/licenses/LICENSE-2.0
9 | #
10 | # Unless required by applicable law or agreed to in writing, software
11 | # distributed under the License is distributed on an "AS IS" BASIS,
12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 | # See the License for the specific language governing permissions and
14 | # limitations under the License.
15 |
16 | # This script converts LLaMA model checkpoint trained by EsayLM to the
17 | # HuggingFace transformers LLaMA PyTorch format, which can then be loaded
18 | # by HuggingFace transformers.
19 |
20 | import gc
21 | import json
22 | import math
23 | import os
24 | import shutil
25 |
26 | import numpy as np
27 | import mlxu
28 | import jax
29 | import jax.numpy as jnp
30 | import flax
31 | from flax.traverse_util import flatten_dict
32 | import torch
33 | from transformers import LlamaConfig, LlamaForCausalLM
34 |
35 | from EasyLM.checkpoint import StreamingCheckpointer
36 | from EasyLM.jax_utils import float_tensor_to_dtype
37 | from EasyLM.models.llama.llama_model import LLAMA_STANDARD_CONFIGS
38 |
39 |
40 | FLAGS, FLAGS_DEF = mlxu.define_flags_with_default(
41 | load_checkpoint='',
42 | model_size='vqlm_1b',
43 | output_dir='',
44 | )
45 |
46 |
47 | def match_keywords(string, positives, negatives):
48 | for positive in positives:
49 | if positive not in string:
50 | return False
51 | for negative in negatives:
52 | if negative in string:
53 | return False
54 | return True
55 |
56 |
57 | def load_and_convert_checkpoint(path):
58 | _, flax_params = StreamingCheckpointer.load_trainstate_checkpoint(path)
59 | flax_params = flatten_dict(flax_params['params'], sep='.')
60 | torch_params = {}
61 | for key, tensor in flax_params.items():
62 | if match_keywords(key, ["kernel"], ["norm", 'ln_f']):
63 | tensor = tensor.T
64 | torch_params[key] = torch.tensor(
65 | float_tensor_to_dtype(tensor, 'fp32'), dtype=torch.float16
66 | )
67 | return torch_params
68 |
69 |
70 | def read_json(path):
71 | with open(path, "r") as f:
72 | return json.load(f)
73 |
74 |
75 | def write_json(text, path):
76 | with open(path, "w") as f:
77 | json.dump(text, f)
78 |
79 |
80 | def write_model(loaded, model_path, model_size):
81 | os.makedirs(model_path, exist_ok=True)
82 | tmp_model_path = os.path.join(model_path, "tmp")
83 | os.makedirs(tmp_model_path, exist_ok=True)
84 |
85 | params = LLAMA_STANDARD_CONFIGS[model_size]
86 |
87 | n_layers = params["num_hidden_layers"]
88 | n_heads = params["num_attention_heads"]
89 | dim = params["hidden_size"]
90 | dims_per_head = dim // n_heads
91 | base = 10000.0
92 | inv_freq = 1.0 / (base ** (torch.arange(0, dims_per_head, 2).float() / dims_per_head))
93 |
94 | # permute for sliced rotary
95 | def permute(w):
96 | return w.view(n_heads, dim // n_heads // 2, 2, dim).transpose(1, 2).reshape(dim, dim)
97 |
98 |
99 | param_count = 0
100 | index_dict = {"weight_map": {}}
101 | for layer_i in range(n_layers):
102 | filename = f"pytorch_model-{layer_i + 1}-of-{n_layers + 1}.bin"
103 | state_dict = {
104 | f"model.layers.{layer_i}.self_attn.q_proj.weight": permute(
105 | loaded[f"transformer.h.{layer_i}.attention.wq.kernel"]
106 | ),
107 | f"model.layers.{layer_i}.self_attn.k_proj.weight": permute(
108 | loaded[f"transformer.h.{layer_i}.attention.wk.kernel"]
109 | ),
110 | f"model.layers.{layer_i}.self_attn.v_proj.weight": loaded[f"transformer.h.{layer_i}.attention.wv.kernel"],
111 | f"model.layers.{layer_i}.self_attn.o_proj.weight": loaded[f"transformer.h.{layer_i}.attention.wo.kernel"],
112 |
113 | f"model.layers.{layer_i}.mlp.gate_proj.weight": loaded[f"transformer.h.{layer_i}.feed_forward.w1.kernel"],
114 | f"model.layers.{layer_i}.mlp.down_proj.weight": loaded[f"transformer.h.{layer_i}.feed_forward.w2.kernel"],
115 | f"model.layers.{layer_i}.mlp.up_proj.weight": loaded[f"transformer.h.{layer_i}.feed_forward.w3.kernel"],
116 |
117 | f"model.layers.{layer_i}.input_layernorm.weight": loaded[f"transformer.h.{layer_i}.attention_norm.kernel"],
118 | f"model.layers.{layer_i}.post_attention_layernorm.weight": loaded[f"transformer.h.{layer_i}.ffn_norm.kernel"],
119 |
120 | }
121 |
122 | state_dict[f"model.layers.{layer_i}.self_attn.rotary_emb.inv_freq"] = inv_freq
123 | for k, v in state_dict.items():
124 | index_dict["weight_map"][k] = filename
125 | param_count += v.numel()
126 | torch.save(state_dict, os.path.join(tmp_model_path, filename))
127 |
128 | filename = f"pytorch_model-{n_layers + 1}-of-{n_layers + 1}.bin"
129 | # Unsharded
130 | state_dict = {
131 | "model.embed_tokens.weight": loaded["transformer.wte.embedding"],
132 | "model.norm.weight": loaded["transformer.ln_f.kernel"],
133 | "lm_head.weight": loaded["lm_head.kernel"],
134 | }
135 |
136 | for k, v in state_dict.items():
137 | index_dict["weight_map"][k] = filename
138 | param_count += v.numel()
139 | torch.save(state_dict, os.path.join(tmp_model_path, filename))
140 |
141 | # Write configs
142 | index_dict["metadata"] = {"total_size": param_count * 2}
143 | write_json(index_dict, os.path.join(tmp_model_path, "pytorch_model.bin.index.json"))
144 |
145 | config = LlamaConfig(
146 | vocab_size=params["vocab_size"],
147 | max_position_embeddings=params["max_sequence_length"],
148 | hidden_size=dim,
149 | intermediate_size=params["intermediate_size"],
150 | num_attention_heads=params["num_attention_heads"],
151 | num_hidden_layers=params["num_hidden_layers"],
152 | rms_norm_eps=params["rms_norm_eps"],
153 | )
154 | config.save_pretrained(tmp_model_path)
155 |
156 | # Make space so we can load the model properly now.
157 | del state_dict
158 | del loaded
159 | gc.collect()
160 |
161 | print("Loading the checkpoint in a Llama model.")
162 | model = LlamaForCausalLM.from_pretrained(tmp_model_path, torch_dtype=torch.float16)
163 | # Avoid saving this as part of the config.
164 | del model.config._name_or_path
165 |
166 | print("Saving in the Transformers format.")
167 | model.save_pretrained(model_path)
168 | shutil.rmtree(tmp_model_path)
169 |
170 |
171 | def main(argv):
172 | assert FLAGS.load_checkpoint != "" and FLAGS.output_dir != ""
173 | assert FLAGS.model_size in LLAMA_STANDARD_CONFIGS
174 |
175 | write_model(
176 | load_and_convert_checkpoint(FLAGS.load_checkpoint),
177 | model_path=FLAGS.output_dir,
178 | model_size=FLAGS.model_size,
179 | )
180 |
181 |
182 | if __name__ == "__main__":
183 | mlxu.run(main)
--------------------------------------------------------------------------------
/EasyLM/models/llama/llama_train.py:
--------------------------------------------------------------------------------
1 | import pprint
2 | from functools import partial
3 |
4 | from tqdm import tqdm, trange
5 | import numpy as np
6 | import mlxu
7 |
8 | import jax
9 | import jax.numpy as jnp
10 | from jax.experimental.pjit import pjit
11 | from jax.sharding import PartitionSpec as PS
12 | from flax.training.train_state import TrainState
13 |
14 | from EasyLM.data import DatasetFactory
15 | from EasyLM.checkpoint import StreamingCheckpointer
16 | from EasyLM.optimizers import OptimizerFactory
17 | from EasyLM.jax_utils import (
18 | JaxRNG, JaxDistributedConfig, next_rng, match_partition_rules,
19 | cross_entropy_loss_and_accuracy, global_norm, get_float_dtype_by_name,
20 | set_random_seed, average_metrics, get_weight_decay_mask,
21 | make_shard_and_gather_fns, with_sharding_constraint,
22 | )
23 | from EasyLM.models.llama.llama_model import (
24 | LLaMAConfig, FlaxLLaMAForCausalLMModule
25 | )
26 |
27 |
28 | FLAGS, FLAGS_DEF = mlxu.define_flags_with_default(
29 | seed=42,
30 | mesh_dim='1,-1,1',
31 | dtype='fp32',
32 | total_steps=10000,
33 | load_llama_config='',
34 | update_llama_config='',
35 | load_checkpoint='',
36 | load_dataset_state='',
37 | load_metadata='',
38 | log_freq=50,
39 | save_model_freq=0,
40 | save_milestone_freq=0,
41 | eval_steps=0,
42 | tokenizer=LLaMAConfig.get_tokenizer_config(),
43 | train_dataset=DatasetFactory.get_default_config(),
44 | eval_dataset=DatasetFactory.get_default_config(),
45 | optimizer=OptimizerFactory.get_default_config(),
46 | checkpointer=StreamingCheckpointer.get_default_config(),
47 | llama=LLaMAConfig.get_default_config(),
48 | logger=mlxu.WandBLogger.get_default_config(),
49 | log_all_worker=False,
50 | jax_distributed=JaxDistributedConfig.get_default_config(),
51 | )
52 |
53 |
54 | def main(argv):
55 | JaxDistributedConfig.initialize(FLAGS.jax_distributed)
56 | variant = mlxu.get_user_flags(FLAGS, FLAGS_DEF)
57 | flags_config_dict = mlxu.user_flags_to_config_dict(FLAGS, FLAGS_DEF)
58 | logger = mlxu.WandBLogger(
59 | config=FLAGS.logger,
60 | variant=variant,
61 | enable=FLAGS.log_all_worker or (jax.process_index() == 0),
62 | )
63 | set_random_seed(FLAGS.seed)
64 |
65 | tokenizer = LLaMAConfig.get_tokenizer(FLAGS.tokenizer)
66 | dataset = DatasetFactory.load_dataset(
67 | FLAGS.train_dataset, tokenizer, device_count=jax.device_count()
68 | )
69 | if FLAGS.load_dataset_state != '':
70 | dataset.load_state_dict(mlxu.load_pickle(FLAGS.load_dataset_state))
71 |
72 | if FLAGS.eval_steps > 0:
73 | eval_dataset = DatasetFactory.load_dataset(
74 | FLAGS.eval_dataset, dataset.tokenizer
75 | )
76 | eval_iterator = iter(eval_dataset)
77 |
78 | seq_length = dataset.seq_length
79 |
80 | if FLAGS.load_llama_config != '':
81 | llama_config = LLaMAConfig.load_config(FLAGS.load_llama_config)
82 | else:
83 | llama_config = LLaMAConfig(**FLAGS.llama)
84 |
85 | if FLAGS.update_llama_config != '':
86 | llama_config.update(dict(eval(FLAGS.update_llama_config)))
87 |
88 | # llama_config.update(dict(
89 | # bos_token_id=dataset.tokenizer.bos_token_id,
90 | # eos_token_id=dataset.tokenizer.eos_token_id,
91 | # ))
92 | # if llama_config.vocab_size < dataset.vocab_size:
93 | # llama_config.update(dict(vocab_size=dataset.vocab_size))
94 |
95 | model = FlaxLLaMAForCausalLMModule(
96 | llama_config, dtype=get_float_dtype_by_name(FLAGS.dtype)
97 | )
98 |
99 | optimizer, optimizer_info = OptimizerFactory.get_optimizer(
100 | FLAGS.optimizer,
101 | get_weight_decay_mask(LLaMAConfig.get_weight_decay_exclusions())
102 | )
103 |
104 | def create_trainstate_from_params(params):
105 | return TrainState.create(params=params, tx=optimizer, apply_fn=None)
106 |
107 | def init_fn(rng):
108 | rng_generator = JaxRNG(rng)
109 | params = model.init(
110 | input_ids=jnp.zeros((4, seq_length), dtype=jnp.int32),
111 | position_ids=jnp.zeros((4, seq_length), dtype=jnp.int32),
112 | attention_mask=jnp.ones((4, seq_length), dtype=jnp.int32),
113 | rngs=rng_generator(llama_config.rng_keys()),
114 | )
115 | return TrainState.create(params=params, tx=optimizer, apply_fn=None)
116 |
117 | def train_step(train_state, rng, batch):
118 | rng_generator = JaxRNG(rng)
119 | batch = with_sharding_constraint(batch, PS(('dp', 'fsdp')))
120 | def loss_and_accuracy(params):
121 | logits = model.apply(
122 | params, batch['input_tokens'], deterministic=False,
123 | rngs=rng_generator(llama_config.rng_keys()),
124 | ).logits
125 | return cross_entropy_loss_and_accuracy(
126 | logits, batch['target_tokens'], batch['loss_masks']
127 | )
128 | grad_fn = jax.value_and_grad(loss_and_accuracy, has_aux=True)
129 | (loss, accuracy), grads = grad_fn(train_state.params)
130 | train_state = train_state.apply_gradients(grads=grads)
131 | metrics = dict(
132 | loss=loss,
133 | accuracy=accuracy,
134 | learning_rate=optimizer_info['learning_rate_schedule'](train_state.step),
135 | gradient_norm=global_norm(grads),
136 | param_norm=global_norm(train_state.params),
137 | )
138 | return train_state, rng_generator(), metrics
139 |
140 | def eval_step(train_state, rng, batch):
141 | rng_generator = JaxRNG(rng)
142 | batch = with_sharding_constraint(batch, PS(('dp', 'fsdp')))
143 | logits = model.apply(
144 | train_state.params, batch['input_tokens'], deterministic=True,
145 | rngs=rng_generator(llama_config.rng_keys()),
146 | ).logits
147 | loss, accuracy = cross_entropy_loss_and_accuracy(
148 | logits, batch['target_tokens'], batch['loss_masks']
149 | )
150 | metrics = dict(
151 | eval_loss=loss,
152 | eval_accuracy=accuracy,
153 | )
154 | return rng_generator(), metrics
155 |
156 | train_state_shapes = jax.eval_shape(init_fn, next_rng())
157 | train_state_partition = match_partition_rules(
158 | LLaMAConfig.get_partition_rules(), train_state_shapes
159 | )
160 |
161 | shard_fns, gather_fns = make_shard_and_gather_fns(
162 | train_state_partition, train_state_shapes
163 | )
164 | checkpointer = StreamingCheckpointer(
165 | FLAGS.checkpointer, logger.output_dir,
166 | enable=jax.process_index() == 0,
167 | )
168 |
169 | sharded_init_fn = pjit(
170 | init_fn,
171 | in_shardings=PS(),
172 | out_shardings=train_state_partition
173 | )
174 |
175 | sharded_create_trainstate_from_params = pjit(
176 | create_trainstate_from_params,
177 | in_shardings=(train_state_partition.params, ),
178 | out_shardings=train_state_partition,
179 | donate_argnums=(0, ),
180 | )
181 |
182 | sharded_train_step = pjit(
183 | train_step,
184 | in_shardings=(train_state_partition, PS(), PS()),
185 | out_shardings=(train_state_partition, PS(), PS()),
186 | donate_argnums=(0, 1),
187 | )
188 |
189 | sharded_eval_step = pjit(
190 | eval_step,
191 | in_shardings=(train_state_partition, PS(), PS()),
192 | out_shardings=(PS(), PS()),
193 | donate_argnums=(1,),
194 | )
195 |
196 | def save_checkpoint(train_state, milestone=False):
197 | step = int(jax.device_get(train_state.step))
198 | metadata = dict(
199 | step=step,
200 | variant=variant,
201 | flags=flags_config_dict,
202 | llama_config=llama_config.to_dict(),
203 | )
204 | checkpointer.save_all(
205 | train_state=train_state,
206 | gather_fns=gather_fns,
207 | metadata=metadata,
208 | dataset=dataset.get_state_dict(),
209 | milestone=milestone,
210 | )
211 |
212 | mesh = LLaMAConfig.get_jax_mesh(FLAGS.mesh_dim)
213 | with mesh:
214 | train_state, restored_params = None, None
215 | if FLAGS.load_checkpoint != '':
216 | train_state, restored_params = checkpointer.load_trainstate_checkpoint(
217 | FLAGS.load_checkpoint, train_state_shapes, shard_fns
218 | )
219 |
220 | if train_state is None and restored_params is None:
221 | # Initialize from scratch
222 | train_state = sharded_init_fn(next_rng())
223 | elif train_state is None and restored_params is not None:
224 | # Restore from params but initialize train_state
225 | train_state = sharded_create_trainstate_from_params(restored_params)
226 | del restored_params
227 | # Hack to get the correct step from the metadata. Otherwise, the `train_state.step` is initialized at 0
228 | if FLAGS.load_metadata != '':
229 | loaded_step = mlxu.load_pickle(FLAGS.load_metadata)["step"] + 1
230 | train_state = train_state.replace(step=loaded_step)
231 | print(f"Resuming at step {loaded_step}")
232 |
233 | start_step = int(jax.device_get(train_state.step))
234 |
235 | # if FLAGS.save_model_freq > 0:
236 | # save_checkpoint(train_state)
237 |
238 | sharded_rng = next_rng()
239 |
240 | step_counter = trange(start_step, FLAGS.total_steps, ncols=0)
241 |
242 | for step, (batch, dataset_metrics) in zip(step_counter, dataset):
243 | train_state, sharded_rng, metrics = sharded_train_step(
244 | train_state, sharded_rng, batch
245 | )
246 |
247 | if step % FLAGS.log_freq == 0:
248 | if FLAGS.eval_steps > 0:
249 | eval_metric_list = []
250 | for _ in range(FLAGS.eval_steps):
251 | eval_batch, _ = next(eval_iterator)
252 | sharded_rng, eval_metrics = sharded_eval_step(
253 | train_state, sharded_rng, eval_batch
254 | )
255 | eval_metric_list.append(eval_metrics)
256 | metrics.update(average_metrics(eval_metric_list))
257 |
258 | log_metrics = {"step": step}
259 | log_metrics.update(metrics)
260 | log_metrics.update(dataset_metrics)
261 | log_metrics = jax.device_get(log_metrics)
262 | logger.log(log_metrics)
263 | if jax.process_index() == 0:
264 | tqdm.write("\n" + pprint.pformat(log_metrics) + "\n")
265 |
266 | if FLAGS.save_milestone_freq > 0 and (step + 1) % FLAGS.save_milestone_freq == 0:
267 | save_checkpoint(train_state, milestone=True)
268 | elif FLAGS.save_model_freq > 0 and (step + 1) % FLAGS.save_model_freq == 0:
269 | save_checkpoint(train_state)
270 |
271 | if FLAGS.save_model_freq > 0:
272 | save_checkpoint(train_state)
273 |
274 |
275 | if __name__ == "__main__":
276 | mlxu.run(main)
277 |
--------------------------------------------------------------------------------
/EasyLM/optimizers.py:
--------------------------------------------------------------------------------
1 | import os
2 | import time
3 | from typing import Any, Mapping, Text, Tuple, Union, NamedTuple
4 | from functools import partial
5 | import re
6 | import dataclasses
7 | import random
8 |
9 | from ml_collections.config_dict import config_dict
10 | from ml_collections import ConfigDict
11 | import jax
12 | import jax.numpy as jnp
13 | import numpy as np
14 | from absl import logging
15 | import optax
16 |
17 | from EasyLM.jax_utils import float_to_dtype
18 |
19 |
20 | class OptimizerFactory(object):
21 | """ Configurable optax optimizer factory. """
22 |
23 | def __init__(self):
24 | raise NotImplementedError
25 |
26 | @staticmethod
27 | def get_default_config(updates=None):
28 | config = ConfigDict()
29 | config.accumulate_gradient_steps = 1
30 | config.type = 'adamw'
31 | config.palm_optimizer = PalmOptimizerFactory.get_default_config()
32 | config.adamw_optimizer = AdamWOptimizerFactory.get_default_config()
33 |
34 | if updates is not None:
35 | config.update(ConfigDict(updates).copy_and_resolve_references())
36 | return config
37 |
38 | @classmethod
39 | def get_optimizer(cls, config, weight_decay_mask=None):
40 | config = cls.get_default_config(config)
41 | if config.type == 'palm':
42 | optimizer, optimizer_info = PalmOptimizerFactory.get_optimizer(
43 | config.palm_optimizer, weight_decay_mask
44 | )
45 | elif config.type == 'adamw':
46 | optimizer, optimizer_info = AdamWOptimizerFactory.get_optimizer(
47 | config.adamw_optimizer, weight_decay_mask
48 | )
49 | else:
50 | raise ValueError(f'Unknown optimizer type: {config.type}')
51 |
52 | if config.accumulate_gradient_steps > 1:
53 | optimizer = optax.MultiSteps(
54 | optimizer, config.accumulate_gradient_steps
55 | )
56 |
57 | return optimizer, optimizer_info
58 |
59 |
60 | class PalmOptimizerFactory(object):
61 | """ PaLM optimizer factory. This optimizer implements the optimizer
62 | described in the PaLM paper: https://arxiv.org/abs/2204.02311
63 | """
64 |
65 | def __init__(self):
66 | raise NotImplementedError
67 |
68 | @staticmethod
69 | def get_default_config(updates=None):
70 | config = ConfigDict()
71 | config.lr = 0.01
72 | config.lr_warmup_steps = 10000
73 | config.b1 = 0.9
74 | config.b2 = 0.99
75 | config.clip_gradient = 1.0
76 | config.weight_decay = 1e-4
77 | config.bf16_momentum = False
78 |
79 | if updates is not None:
80 | config.update(ConfigDict(updates).copy_and_resolve_references())
81 | return config
82 |
83 | @classmethod
84 | def get_optimizer(cls, config, weight_decay_mask=None):
85 | config = cls.get_default_config(config)
86 |
87 | def learning_rate_schedule(step):
88 | multiplier = config.lr / 0.01
89 | return multiplier / jnp.sqrt(jnp.maximum(step, config.lr_warmup_steps))
90 |
91 | def weight_decay_schedule(step):
92 | multiplier = config.weight_decay / 1e-4
93 | return -multiplier * jnp.square(learning_rate_schedule(step))
94 |
95 | optimizer_info = dict(
96 | learning_rate_schedule=learning_rate_schedule,
97 | weight_decay_schedule=weight_decay_schedule,
98 | )
99 |
100 | optimizer = optax.chain(
101 | optax.clip_by_global_norm(config.clip_gradient),
102 | optax.adafactor(
103 | learning_rate=learning_rate_schedule,
104 | multiply_by_parameter_scale=True,
105 | momentum=config.b1,
106 | decay_rate=config.b2,
107 | factored=False,
108 | clipping_threshold=None,
109 | dtype_momentum=jnp.bfloat16 if config.bf16_momentum else jnp.float32,
110 | ),
111 | optax_add_scheduled_weight_decay(
112 | weight_decay_schedule, weight_decay_mask
113 | )
114 | )
115 | return optimizer, optimizer_info
116 |
117 |
118 | class AdamWOptimizerFactory(object):
119 | """ AdamW optimizer with cosine schedule. """
120 |
121 | def __init__(self):
122 | raise NotImplementedError
123 |
124 | @staticmethod
125 | def get_default_config(updates=None):
126 | config = ConfigDict()
127 | config.init_lr = 0.0
128 | config.end_lr = 0.001
129 | config.lr = 0.01
130 | config.lr_warmup_steps = 2000
131 | config.lr_decay_steps = 500000
132 | config.b1 = 0.9
133 | config.b2 = 0.95
134 | config.clip_gradient = 1.0
135 | config.weight_decay = 1e-4
136 | config.bf16_momentum = False
137 | config.multiply_by_parameter_scale = False
138 |
139 | if updates is not None:
140 | config.update(ConfigDict(updates).copy_and_resolve_references())
141 | return config
142 |
143 | @classmethod
144 | def get_optimizer(cls, config, weight_decay_mask=None):
145 | config = cls.get_default_config(config)
146 |
147 | learning_rate_schedule = optax.warmup_cosine_decay_schedule(
148 | init_value=config.init_lr,
149 | peak_value=config.lr,
150 | warmup_steps=config.lr_warmup_steps,
151 | decay_steps=config.lr_decay_steps,
152 | end_value=config.end_lr,
153 | )
154 |
155 | optimizer_info = dict(
156 | learning_rate_schedule=learning_rate_schedule,
157 | )
158 |
159 | if config.multiply_by_parameter_scale:
160 | optimizer = optax.chain(
161 | optax.clip_by_global_norm(config.clip_gradient),
162 | optax.adafactor(
163 | learning_rate=learning_rate_schedule,
164 | multiply_by_parameter_scale=True,
165 | momentum=config.b1,
166 | decay_rate=config.b2,
167 | factored=False,
168 | clipping_threshold=None,
169 | dtype_momentum=jnp.bfloat16 if config.bf16_momentum else jnp.float32,
170 | ),
171 | optax_add_scheduled_weight_decay(
172 | lambda step: -learning_rate_schedule(step) * config.weight_decay,
173 | weight_decay_mask
174 | )
175 | )
176 | else:
177 | optimizer = optax.chain(
178 | optax.clip_by_global_norm(config.clip_gradient),
179 | optax.adamw(
180 | learning_rate=learning_rate_schedule,
181 | weight_decay=config.weight_decay,
182 | b1=config.b1,
183 | b2=config.b2,
184 | mask=weight_decay_mask,
185 | mu_dtype=jnp.bfloat16 if config.bf16_momentum else jnp.float32,
186 | ),
187 | )
188 |
189 | return optimizer, optimizer_info
190 |
191 |
192 | class OptaxScheduledWeightDecayState(NamedTuple):
193 | count: jax.Array
194 |
195 |
196 | def optax_add_scheduled_weight_decay(schedule_fn, mask=None):
197 | """ Apply weight decay with schedule. """
198 |
199 | def init_fn(params):
200 | del params
201 | return OptaxScheduledWeightDecayState(count=jnp.zeros([], jnp.int32))
202 |
203 | def update_fn(updates, state, params):
204 | if params is None:
205 | raise ValueError('Params cannot be None for weight decay!')
206 |
207 | weight_decay = schedule_fn(state.count)
208 | updates = jax.tree_util.tree_map(
209 | lambda g, p: g + weight_decay * p, updates, params
210 | )
211 | return updates, OptaxScheduledWeightDecayState(
212 | count=optax.safe_int32_increment(state.count)
213 | )
214 |
215 | if mask is not None:
216 | return optax.masked(optax.GradientTransformation(init_fn, update_fn), mask)
217 | return optax.GradientTransformation(init_fn, update_fn)
218 |
--------------------------------------------------------------------------------
/EasyLM/scripts/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/ytongbai/LVM/b6de939ef0eb1ee6593445a7f5268145f338749b/EasyLM/scripts/__init__.py
--------------------------------------------------------------------------------
/EasyLM/scripts/convert_checkpoint.py:
--------------------------------------------------------------------------------
1 | # This script converts model checkpoint trained by EsayLM to a standard
2 | # mspack checkpoint that can be loaded by huggingface transformers or
3 | # flax.serialization.msgpack_restore. Such conversion allows models to be
4 | # used by other frameworks that integrate with huggingface transformers.
5 |
6 | import pprint
7 | from functools import partial
8 | import os
9 | import numpy as np
10 | import mlxu
11 | import jax.numpy as jnp
12 | import flax.serialization
13 | from EasyLM.checkpoint import StreamingCheckpointer
14 | from EasyLM.jax_utils import float_to_dtype
15 |
16 |
17 | FLAGS, FLAGS_DEF = mlxu.define_flags_with_default(
18 | load_checkpoint='',
19 | output_file='',
20 | streaming=False,
21 | float_dtype='bf16',
22 | )
23 |
24 |
25 | def main(argv):
26 | assert FLAGS.load_checkpoint != '' and FLAGS.output_file != '', 'input and output must be specified'
27 | params = StreamingCheckpointer.load_trainstate_checkpoint(
28 | FLAGS.load_checkpoint, disallow_trainstate=True
29 | )[1]['params']
30 |
31 | if FLAGS.streaming:
32 | StreamingCheckpointer.save_train_state_to_file(
33 | params, FLAGS.output_file, float_dtype=FLAGS.float_dtype
34 | )
35 | else:
36 | params = float_to_dtype(params, FLAGS.float_dtype)
37 | with mlxu.open_file(FLAGS.output, 'wb') as fout:
38 | fout.write(flax.serialization.msgpack_serialize(params, in_place=True))
39 |
40 |
41 | if __name__ == "__main__":
42 | mlxu.run(main)
43 |
--------------------------------------------------------------------------------
/EasyLM/scripts/diff_checkpoint.py:
--------------------------------------------------------------------------------
1 | # This script converts model checkpoint trained by EsayLM to a standard
2 | # mspack checkpoint that can be loaded by huggingface transformers or
3 | # flax.serialization.msgpack_restore. Such conversion allows models to be
4 | # used by other frameworks that integrate with huggingface transformers.
5 |
6 | import pprint
7 | from functools import partial
8 | import os
9 | import numpy as np
10 | import jax
11 | import jax.numpy as jnp
12 | import flax.serialization
13 | import mlxu
14 | from EasyLM.checkpoint import StreamingCheckpointer
15 | from EasyLM.jax_utils import float_to_dtype
16 |
17 |
18 | FLAGS, FLAGS_DEF = mlxu.define_flags_with_default(
19 | recover_diff=False,
20 | load_base_checkpoint='',
21 | load_target_checkpoint='',
22 | output_file='',
23 | streaming=True,
24 | float_dtype='bf16',
25 | )
26 |
27 |
28 | def main(argv):
29 | assert FLAGS.load_base_checkpoint != '' and FLAGS.load_target_checkpoint != ''
30 | assert FLAGS.output_file != ''
31 | base_params = StreamingCheckpointer.load_trainstate_checkpoint(
32 | FLAGS.load_base_checkpoint, disallow_trainstate=True
33 | )[1]['params']
34 |
35 | target_params = StreamingCheckpointer.load_trainstate_checkpoint(
36 | FLAGS.load_target_checkpoint, disallow_trainstate=True
37 | )[1]['params']
38 |
39 | if FLAGS.recover_diff:
40 | params = jax.tree_util.tree_map(
41 | lambda b, t: b + t, base_params, target_params
42 | )
43 | else:
44 | params = jax.tree_util.tree_map(
45 | lambda b, t: t - b, base_params, target_params
46 | )
47 |
48 | if FLAGS.streaming:
49 | StreamingCheckpointer.save_train_state_to_file(
50 | params, FLAGS.output_file, float_dtype=FLAGS.float_dtype
51 | )
52 | else:
53 | params = float_to_dtype(params, FLAGS.float_dtype)
54 | with mlxu.open_file(FLAGS.output, 'wb') as fout:
55 | fout.write(flax.serialization.msgpack_serialize(params, in_place=True))
56 |
57 |
58 | if __name__ == "__main__":
59 | mlxu.run(main)
60 |
--------------------------------------------------------------------------------
/EasyLM/scripts/lm_eval_harness.py:
--------------------------------------------------------------------------------
1 | # This script runs lm_eval_harness evaluations against a served language model.
2 | # Typically, you need to run a language model server first, e.g.:
3 | # python -m EasyLM.models.gptj.gptj_serve ...
4 |
5 | import dataclasses
6 | import pprint
7 | from functools import partial
8 | import os
9 | from tqdm import tqdm, trange
10 | import numpy as np
11 | import mlxu
12 |
13 | from flax.traverse_util import flatten_dict
14 | from lm_eval import evaluator, tasks
15 | from lm_eval.base import LM
16 |
17 | from EasyLM.serving import LMClient
18 |
19 |
20 | FLAGS, FLAGS_DEF = mlxu.define_flags_with_default(
21 | tasks='wsc,piqa,winogrande,openbookqa,logiqa',
22 | shots=0,
23 | limit=0,
24 | write_out=False,
25 | lm_client=LMClient.get_default_config(),
26 | logger=mlxu.WandBLogger.get_default_config(),
27 | )
28 |
29 |
30 | class LMEvalHarnessInterface(LM):
31 |
32 | def __init__(self, lm_client):
33 | self.lm_client = lm_client
34 |
35 | def greedy_until(self, inputs):
36 | prefix, until = zip(*inputs)
37 | return self.lm_client.greedy_until(prefix, until)
38 |
39 | def loglikelihood_rolling(self, inputs):
40 | loglikelihood, is_greedy = self.lm_client.loglikelihood_rolling(inputs)
41 | return list(zip(loglikelihood, is_greedy))
42 |
43 | def loglikelihood(self, inputs):
44 | prefix, text = zip(*inputs)
45 | loglikelihood, is_greedy = self.lm_client.loglikelihood(prefix, text)
46 | return list(zip(loglikelihood, is_greedy))
47 |
48 |
49 | def main(argv):
50 | logger = mlxu.WandBLogger(
51 | config=FLAGS.logger, variant=mlxu.get_user_flags(FLAGS, FLAGS_DEF)
52 | )
53 | model = LMEvalHarnessInterface(LMClient(FLAGS.lm_client))
54 | task_list = FLAGS.tasks.split(',')
55 | results = evaluator.evaluate(
56 | model, tasks.get_task_dict(task_list), False, FLAGS.shots,
57 | limit=None if FLAGS.limit <= 0 else FLAGS.limit,
58 | write_out=FLAGS.write_out,
59 | )
60 | logger.log(flatten_dict(results['results'], sep='/'))
61 | pprint.pprint(results)
62 |
63 |
64 | if __name__ == "__main__":
65 | mlxu.run(main)
66 |
--------------------------------------------------------------------------------
/EasyLM/scripts/lm_eval_json.py:
--------------------------------------------------------------------------------
1 | import json
2 | import mlxu
3 | from EasyLM.serving import LMClient
4 |
5 |
6 | FLAGS, FLAGS_DEF = mlxu.define_flags_with_default(
7 | input_file='',
8 | output_file='',
9 | prefix_field='prefix',
10 | text_field='text',
11 | until_field='until',
12 | eval_type='loglikelihood',
13 | lm_client=LMClient.get_default_config(),
14 | )
15 |
16 |
17 | def main(argv):
18 | lm_client = LMClient(FLAGS.lm_client)
19 | with mlxu.open_file(FLAGS.input_file, 'r') as fin:
20 | input_data = json.load(fin)
21 |
22 | if FLAGS.eval_type == 'loglikelihood':
23 | prefix = input_data[FLAGS.prefix_field]
24 | text = input_data[FLAGS.text_field]
25 | loglikelihoods, is_greedys = lm_client.loglikelihood(prefix, text)
26 | output_data = {
27 | 'loglikelihood': loglikelihoods,
28 | 'is_greedy': is_greedys,
29 | }
30 | elif FLAGS.eval_type == 'loglikelihood_rolling':
31 | text = input_data[FLAGS.text_field]
32 | loglikelihoods, is_greedys = lm_client.loglikelihood_rolling(text)
33 | output_data = {
34 | 'loglikelihood': loglikelihoods,
35 | 'is_greedy': is_greedys,
36 | }
37 | elif FLAGS.eval_type == 'greedy_until':
38 | prefix = input_data[FLAGS.prefix_field]
39 | until = input_data[FLAGS.until_field]
40 | output_data = {'output_text': lm_client.greedy_until(prefix, until)}
41 | elif FLAGS.eval_type == 'generate':
42 | prefix = input_data[FLAGS.prefix_field]
43 | output_data = {'output_text': lm_client.generate(prefix)}
44 | else:
45 | raise ValueError(f'Unknown eval_type: {FLAGS.eval_type}')
46 |
47 | with mlxu.open_file(FLAGS.output_file, 'w') as fout:
48 | json.dump(output_data, fout)
49 |
50 |
51 | if __name__ == "__main__":
52 | mlxu.run(main)
53 |
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 | Apache License
2 | Version 2.0, January 2004
3 | http://www.apache.org/licenses/
4 |
5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6 |
7 | 1. Definitions.
8 |
9 | "License" shall mean the terms and conditions for use, reproduction,
10 | and distribution as defined by Sections 1 through 9 of this document.
11 |
12 | "Licensor" shall mean the copyright owner or entity authorized by
13 | the copyright owner that is granting the License.
14 |
15 | "Legal Entity" shall mean the union of the acting entity and all
16 | other entities that control, are controlled by, or are under common
17 | control with that entity. For the purposes of this definition,
18 | "control" means (i) the power, direct or indirect, to cause the
19 | direction or management of such entity, whether by contract or
20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the
21 | outstanding shares, or (iii) beneficial ownership of such entity.
22 |
23 | "You" (or "Your") shall mean an individual or Legal Entity
24 | exercising permissions granted by this License.
25 |
26 | "Source" form shall mean the preferred form for making modifications,
27 | including but not limited to software source code, documentation
28 | source, and configuration files.
29 |
30 | "Object" form shall mean any form resulting from mechanical
31 | transformation or translation of a Source form, including but
32 | not limited to compiled object code, generated documentation,
33 | and conversions to other media types.
34 |
35 | "Work" shall mean the work of authorship, whether in Source or
36 | Object form, made available under the License, as indicated by a
37 | copyright notice that is included in or attached to the work
38 | (an example is provided in the Appendix below).
39 |
40 | "Derivative Works" shall mean any work, whether in Source or Object
41 | form, that is based on (or derived from) the Work and for which the
42 | editorial revisions, annotations, elaborations, or other modifications
43 | represent, as a whole, an original work of authorship. For the purposes
44 | of this License, Derivative Works shall not include works that remain
45 | separable from, or merely link (or bind by name) to the interfaces of,
46 | the Work and Derivative Works thereof.
47 |
48 | "Contribution" shall mean any work of authorship, including
49 | the original version of the Work and any modifications or additions
50 | to that Work or Derivative Works thereof, that is intentionally
51 | submitted to Licensor for inclusion in the Work by the copyright owner
52 | or by an individual or Legal Entity authorized to submit on behalf of
53 | the copyright owner. For the purposes of this definition, "submitted"
54 | means any form of electronic, verbal, or written communication sent
55 | to the Licensor or its representatives, including but not limited to
56 | communication on electronic mailing lists, source code control systems,
57 | and issue tracking systems that are managed by, or on behalf of, the
58 | Licensor for the purpose of discussing and improving the Work, but
59 | excluding communication that is conspicuously marked or otherwise
60 | designated in writing by the copyright owner as "Not a Contribution."
61 |
62 | "Contributor" shall mean Licensor and any individual or Legal Entity
63 | on behalf of whom a Contribution has been received by Licensor and
64 | subsequently incorporated within the Work.
65 |
66 | 2. Grant of Copyright License. Subject to the terms and conditions of
67 | this License, each Contributor hereby grants to You a perpetual,
68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable
69 | copyright license to reproduce, prepare Derivative Works of,
70 | publicly display, publicly perform, sublicense, and distribute the
71 | Work and such Derivative Works in Source or Object form.
72 |
73 | 3. Grant of Patent License. Subject to the terms and conditions of
74 | this License, each Contributor hereby grants to You a perpetual,
75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable
76 | (except as stated in this section) patent license to make, have made,
77 | use, offer to sell, sell, import, and otherwise transfer the Work,
78 | where such license applies only to those patent claims licensable
79 | by such Contributor that are necessarily infringed by their
80 | Contribution(s) alone or by combination of their Contribution(s)
81 | with the Work to which such Contribution(s) was submitted. If You
82 | institute patent litigation against any entity (including a
83 | cross-claim or counterclaim in a lawsuit) alleging that the Work
84 | or a Contribution incorporated within the Work constitutes direct
85 | or contributory patent infringement, then any patent licenses
86 | granted to You under this License for that Work shall terminate
87 | as of the date such litigation is filed.
88 |
89 | 4. Redistribution. You may reproduce and distribute copies of the
90 | Work or Derivative Works thereof in any medium, with or without
91 | modifications, and in Source or Object form, provided that You
92 | meet the following conditions:
93 |
94 | (a) You must give any other recipients of the Work or
95 | Derivative Works a copy of this License; and
96 |
97 | (b) You must cause any modified files to carry prominent notices
98 | stating that You changed the files; and
99 |
100 | (c) You must retain, in the Source form of any Derivative Works
101 | that You distribute, all copyright, patent, trademark, and
102 | attribution notices from the Source form of the Work,
103 | excluding those notices that do not pertain to any part of
104 | the Derivative Works; and
105 |
106 | (d) If the Work includes a "NOTICE" text file as part of its
107 | distribution, then any Derivative Works that You distribute must
108 | include a readable copy of the attribution notices contained
109 | within such NOTICE file, excluding those notices that do not
110 | pertain to any part of the Derivative Works, in at least one
111 | of the following places: within a NOTICE text file distributed
112 | as part of the Derivative Works; within the Source form or
113 | documentation, if provided along with the Derivative Works; or,
114 | within a display generated by the Derivative Works, if and
115 | wherever such third-party notices normally appear. The contents
116 | of the NOTICE file are for informational purposes only and
117 | do not modify the License. You may add Your own attribution
118 | notices within Derivative Works that You distribute, alongside
119 | or as an addendum to the NOTICE text from the Work, provided
120 | that such additional attribution notices cannot be construed
121 | as modifying the License.
122 |
123 | You may add Your own copyright statement to Your modifications and
124 | may provide additional or different license terms and conditions
125 | for use, reproduction, or distribution of Your modifications, or
126 | for any such Derivative Works as a whole, provided Your use,
127 | reproduction, and distribution of the Work otherwise complies with
128 | the conditions stated in this License.
129 |
130 | 5. Submission of Contributions. Unless You explicitly state otherwise,
131 | any Contribution intentionally submitted for inclusion in the Work
132 | by You to the Licensor shall be under the terms and conditions of
133 | this License, without any additional terms or conditions.
134 | Notwithstanding the above, nothing herein shall supersede or modify
135 | the terms of any separate license agreement you may have executed
136 | with Licensor regarding such Contributions.
137 |
138 | 6. Trademarks. This License does not grant permission to use the trade
139 | names, trademarks, service marks, or product names of the Licensor,
140 | except as required for reasonable and customary use in describing the
141 | origin of the Work and reproducing the content of the NOTICE file.
142 |
143 | 7. Disclaimer of Warranty. Unless required by applicable law or
144 | agreed to in writing, Licensor provides the Work (and each
145 | Contributor provides its Contributions) on an "AS IS" BASIS,
146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147 | implied, including, without limitation, any warranties or conditions
148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149 | PARTICULAR PURPOSE. You are solely responsible for determining the
150 | appropriateness of using or redistributing the Work and assume any
151 | risks associated with Your exercise of permissions under this License.
152 |
153 | 8. Limitation of Liability. In no event and under no legal theory,
154 | whether in tort (including negligence), contract, or otherwise,
155 | unless required by applicable law (such as deliberate and grossly
156 | negligent acts) or agreed to in writing, shall any Contributor be
157 | liable to You for damages, including any direct, indirect, special,
158 | incidental, or consequential damages of any character arising as a
159 | result of this License or out of the use or inability to use the
160 | Work (including but not limited to damages for loss of goodwill,
161 | work stoppage, computer failure or malfunction, or any and all
162 | other commercial damages or losses), even if such Contributor
163 | has been advised of the possibility of such damages.
164 |
165 | 9. Accepting Warranty or Additional Liability. While redistributing
166 | the Work or Derivative Works thereof, You may choose to offer,
167 | and charge a fee for, acceptance of support, warranty, indemnity,
168 | or other liability obligations and/or rights consistent with this
169 | License. However, in accepting such obligations, You may act only
170 | on Your own behalf and on Your sole responsibility, not on behalf
171 | of any other Contributor, and only if You agree to indemnify,
172 | defend, and hold each Contributor harmless for any liability
173 | incurred by, or claims asserted against, such Contributor by reason
174 | of your accepting any such warranty or additional liability.
175 |
176 | END OF TERMS AND CONDITIONS
177 |
178 | APPENDIX: How to apply the Apache License to your work.
179 |
180 | To apply the Apache License to your work, attach the following
181 | boilerplate notice, with the fields enclosed by brackets "[]"
182 | replaced with your own identifying information. (Don't include
183 | the brackets!) The text should be enclosed in the appropriate
184 | comment syntax for the file format. We also recommend that a
185 | file or class name and description of purpose be included on the
186 | same "printed page" as the copyright notice for easier
187 | identification within third-party archives.
188 |
189 | Copyright [yyyy] [name of copyright owner]
190 |
191 | Licensed under the Apache License, Version 2.0 (the "License");
192 | you may not use this file except in compliance with the License.
193 | You may obtain a copy of the License at
194 |
195 | http://www.apache.org/licenses/LICENSE-2.0
196 |
197 | Unless required by applicable law or agreed to in writing, software
198 | distributed under the License is distributed on an "AS IS" BASIS,
199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200 | See the License for the specific language governing permissions and
201 | limitations under the License.
202 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 |
2 | # LVM: Sequential Modeling Enables Scalable Learning for Large Vision Models
3 |
4 | [LVM](https://arxiv.org/abs/2312.00785) is a vision pretraining model that converts various kinds of visual data into visual sentences and performs next-token prediction autoregressively. It is compatible with both GPU and TPU.
5 |
6 | LVM is built on top of [OpenLLaMA](https://github.com/openlm-research/open_llama) (an autoregressive model) and [OpenMuse](https://github.com/huggingface/open-muse) (a VQGAN that converts images into visual tokens).
7 |
8 | This was trained in collaboration with HuggingFace. Thanks [Victor Sanh](https://github.com/VictorSanh) for the support in this project.
9 |
10 | ## Abstract:
11 |
12 | We introduce a novel sequential modeling approach which enables learning a Large Vision Model (LVM) without making use of any linguistic data.
13 | To do this, we define a common format, ``visual sentences", in which we can represent raw images and videos as well as annotated data sources such as semantic segmentations and depth reconstructions without needing any meta-knowledge beyond the pixels. Once this wide variety of visual data (comprising 420 billion tokens) is represented as sequences, the model can be trained to minimize a cross-entropy loss for next token prediction. By training across various scales of model architecture and data diversity, we provide empirical evidence that our models scale effectively. Many different vision tasks can be solved by designing suitable visual prompts at test time.
14 |
15 | ## Visual Sentence
16 |
17 |
18 |
19 |