test_llama.py 24.8 KB
Newer Older
Tri Dao's avatar
Tri Dao committed
1
2
3
4
# Copyright (c) 2023, Tri Dao.

# To run the huggingface implementation, we first need to convert the weights:
# https://github.com/huggingface/transformers/pull/21955
5
# python -m transformers.models.llama.convert_llama_weights_to_hf --input_dir $CHECKPOINT_DIR/llama --model_size 7B --output_dir $CHECKPOINT_DIR/llama/7B-hf
Tri Dao's avatar
Tri Dao committed
6
7
8
9
10
# and repeat for 13B, 30B, 65B

import os
import time
from pathlib import Path
11

Tri Dao's avatar
Tri Dao committed
12
13
14
15
current_dir = Path(__file__).parent.absolute()

import torch
import pytest
16
import shutil
Tri Dao's avatar
Tri Dao committed
17

Tri Dao's avatar
Tri Dao committed
18
19
from einops import rearrange

20
from transformers import LlamaTokenizer, LlamaConfig
Tri Dao's avatar
Tri Dao committed
21
22
from transformers.models.llama.modeling_llama import LlamaForCausalLM

Tri Dao's avatar
Tri Dao committed
23
from flash_attn.models.gpt import GPTLMHeadModel, combine_state_dicts_tp, shard_state_dict_tp
24
25
26
27
28
29
from flash_attn.models.llama import (
    remap_state_dict_meta_llama,
    llama_config_to_gpt2_config,
    remap_state_dict_hf_llama,
    inv_remap_state_dict_hf_llama,
)
Tri Dao's avatar
Tri Dao committed
30
from flash_attn.models.llama import config_from_checkpoint, state_dicts_from_checkpoint
Tri Dao's avatar
Tri Dao committed
31
from flash_attn.utils.distributed import all_gather_raw
Tri Dao's avatar
Tri Dao committed
32
33
34
35
from flash_attn.utils.pretrained import state_dict_from_pretrained
from flash_attn.utils.generation import update_graph_cache


36
37
38
39
40
41
def _pretrained_state_dict_from_checkpoint(checkpoint_path, model_name, config, checkpoint_format):
    if checkpoint_format == "meta":
        ckpt_state_dicts = state_dicts_from_checkpoint(checkpoint_path, model_name)
        pretrained_state_dicts = [remap_state_dict_meta_llama(s, config) for s in ckpt_state_dicts]
        pretrained_state_dict = combine_state_dicts_tp(pretrained_state_dicts, config)
    else:
42
43
44
        pretrained_state_dict = state_dict_from_pretrained(
            Path(checkpoint_path) / f"{model_name}-hf"
        )
45
46
47
48
        pretrained_state_dict = remap_state_dict_hf_llama(pretrained_state_dict, config)
    return pretrained_state_dict


49
@pytest.mark.parametrize("model_name", ["7B"])
Tri Dao's avatar
Tri Dao committed
50
def test_llama_state_dict(model_name):
51
52
53
    checkpoint_path = (
        Path(os.environ.get("CHECKPOINT_DIR", current_dir.parent.parent / "checkpoints")) / "llama"
    )
Tri Dao's avatar
Tri Dao committed
54
55
56
    config = llama_config_to_gpt2_config(config_from_checkpoint(checkpoint_path, model_name))
    ckpt_state_dicts = state_dicts_from_checkpoint(checkpoint_path, model_name)
    pretrained_state_dict = remap_state_dict_meta_llama(ckpt_state_dicts[0], config)
57
    model = GPTLMHeadModel(config, device="meta")  # Without device='meta' init is very slow
Tri Dao's avatar
Tri Dao committed
58
    state_dict = model.state_dict()
59
60
    assert state_dict.keys() == pretrained_state_dict.keys()
    for k in state_dict.keys():
Tri Dao's avatar
Tri Dao committed
61
62
63
        assert state_dict[k].shape == pretrained_state_dict[k].shape


64
65
@pytest.mark.parametrize("model_name", ["7B", "13B"])
@pytest.mark.parametrize("checkpoint_format", ["meta", "hf"])
66
def test_llama_optimized(model_name, checkpoint_format):
Tri Dao's avatar
Tri Dao committed
67
68
69
70
    """Check that our implementation of LLaMa (with all optimizations enabled) matches the
    HF implementation: the output of our forward pass in fp16 should be around the same as the HF
    forward pass in fp16, when compared to the HF forward pass in fp32.
    """
71
72
73
    checkpoint_path = (
        Path(os.environ.get("CHECKPOINT_DIR", current_dir.parent.parent / "checkpoints")) / "llama"
    )
Tri Dao's avatar
Tri Dao committed
74
75

    dtype = torch.float16
76
    device = "cuda"
77
78
    config = config_from_checkpoint(checkpoint_path, model_name, checkpoint_format)
    config = llama_config_to_gpt2_config(config)
Tri Dao's avatar
Tri Dao committed
79
80
81
82
83
84
    config.use_flash_attn = True
    config.fused_bias_fc = True
    config.fused_mlp = False  # We don't have fused GatedMLP yet
    config.fused_dropout_add_ln = True
    config.residual_in_fp32 = True

85
86
87
    pretrained_state_dict = _pretrained_state_dict_from_checkpoint(
        checkpoint_path, model_name, config, checkpoint_format
    )
Tri Dao's avatar
Tri Dao committed
88
    model = GPTLMHeadModel(config, device=device, dtype=dtype)
Tri Dao's avatar
Tri Dao committed
89
    model.load_state_dict(pretrained_state_dict)
Tri Dao's avatar
Tri Dao committed
90
91
92
93
94
95
    model.eval()

    torch.manual_seed(0)
    batch_size = 2
    max_seqlen = 256
    seqlens = torch.randint(max_seqlen // 2, max_seqlen + 1, (batch_size,), device=device)
96
97
98
    input_ids = torch.randint(
        0, config.vocab_size, (batch_size, max_seqlen), dtype=torch.long, device=device
    )
Tri Dao's avatar
Tri Dao committed
99
100
101
102
103
104
105
    with torch.no_grad():
        out = model.transformer(input_ids)
        logits = model(input_ids).logits
    del model

    # Without device_map, the model is loaded on the CPU, which is very slow
    # Need auto here since the 13B fp32 model doesn't fit in memory on a A100 40GB
106
107
108
    model_ref = LlamaForCausalLM.from_pretrained(
        Path(checkpoint_path) / f"{model_name}-hf", device_map="auto"
    )
Tri Dao's avatar
Tri Dao committed
109
110
111
112
113
114
    model_ref.eval()
    with torch.no_grad():
        out_ref = model_ref.model(input_ids).last_hidden_state.to(device=device)
        logits_ref = model_ref(input_ids).logits.to(device=device)
    del model_ref

115
116
117
    model_hf = LlamaForCausalLM.from_pretrained(
        Path(checkpoint_path) / f"{model_name}-hf", torch_dtype=dtype, device_map={"": device}
    )
Tri Dao's avatar
Tri Dao committed
118
    model_hf.eval()
Tri Dao's avatar
Tri Dao committed
119
120
121
    with torch.no_grad():
        out_hf = model_hf.model(input_ids).last_hidden_state
        logits_hf = model_hf(input_ids).logits
Tri Dao's avatar
Tri Dao committed
122
123
    del model_hf

124
125
126
127
    print(f"Output max diff: {(out - out_ref).abs().max().item()}")
    print(f"Output mean diff: {(out - out_ref).abs().mean().item()}")
    print(f"HF fp16 max diff: {(out_hf - out_ref).abs().max().item()}")
    print(f"HF fp16 mean diff: {(out_hf - out_ref).abs().mean().item()}")
Tri Dao's avatar
Tri Dao committed
128
129
    assert (out - out_ref).abs().max().item() < 3 * (out_hf - out_ref).abs().max().item()

130
131
132
133
134
135
136
    print(f"Logits max diff: {(logits - logits_ref).abs().max().item()}")
    print(f"Logits mean diff: {(logits - logits_ref).abs().mean().item()}")
    print(f"HF fp16 max diff: {(logits_hf - logits_ref).abs().max().item()}")
    print(f"HF fp16 mean diff: {(logits_hf - logits_ref).abs().mean().item()}")
    assert (logits - logits_ref).abs().max().item() < 3 * (
        logits_hf - logits_ref
    ).abs().max().item()
Tri Dao's avatar
Tri Dao committed
137
138
139


# torchrun --no_python --nproc_per_node=2 pytest -q -s tests/models/test_llama.py -k "parallel"
140
141
142
@pytest.mark.parametrize("world_size", [2])
@pytest.mark.parametrize("model_name", ["13B"])
@pytest.mark.parametrize("checkpoint_format", ["meta", "hf"])
143
def test_llama_parallel(model_name, world_size, checkpoint_format):
Tri Dao's avatar
Tri Dao committed
144
145
146
147
148
149
    """Check that our implementation of LLaMa (with all optimizations enabled) matches the
    HF implementation: the output of our forward pass in fp16 should be around the same as the HF
    forward pass in fp16, when compared to the HF forward pass in fp32.
    """
    from apex.transformer import parallel_state

150
151
152
    checkpoint_path = (
        Path(os.environ.get("CHECKPOINT_DIR", current_dir.parent.parent / "checkpoints")) / "llama"
    )
Tri Dao's avatar
Tri Dao committed
153
154

    dtype = torch.float16
155
156
    config = config_from_checkpoint(checkpoint_path, model_name, checkpoint_format)
    config = llama_config_to_gpt2_config(config)
Tri Dao's avatar
Tri Dao committed
157
158
159
160
161
162
163
    config.use_flash_attn = True
    config.fused_bias_fc = True
    config.fused_mlp = False  # We don't have fused GatedMLP yet
    config.fused_dropout_add_ln = True
    config.residual_in_fp32 = True

    if not torch.distributed.is_initialized():
164
165
        torch.distributed.init_process_group(backend="nccl", init_method="env://")
    device = f"cuda:{torch.distributed.get_rank()}"
Tri Dao's avatar
Tri Dao committed
166
167
168
169
170
    assert world_size <= torch.distributed.get_world_size()
    parallel_state.initialize_model_parallel(tensor_model_parallel_size_=world_size)
    rank = parallel_state.get_tensor_model_parallel_rank()
    process_group = parallel_state.get_tensor_model_parallel_group()

171
172
173
    pretrained_state_dict = _pretrained_state_dict_from_checkpoint(
        checkpoint_path, model_name, config, checkpoint_format
    )
Tri Dao's avatar
Tri Dao committed
174
    model = GPTLMHeadModel(config, process_group=process_group, device=device, dtype=dtype)
Tri Dao's avatar
Tri Dao committed
175
    model.load_state_dict(shard_state_dict_tp(pretrained_state_dict, config, world_size, rank))
Tri Dao's avatar
Tri Dao committed
176
177
178
179
180
181
    model.eval()

    torch.manual_seed(0)
    batch_size = 2
    max_seqlen = 256
    seqlens = torch.randint(max_seqlen // 2, max_seqlen + 1, (batch_size,), device=device)
182
183
184
    input_ids = torch.randint(
        0, config.vocab_size, (batch_size, max_seqlen), dtype=torch.long, device=device
    )
Tri Dao's avatar
Tri Dao committed
185
186
    with torch.no_grad():
        out = model.transformer(input_ids)
Tri Dao's avatar
Tri Dao committed
187
188
        out, _ = all_gather_raw(out, process_group=process_group)
        out = rearrange(out, "(b s) d -> b s d", b=batch_size)
Tri Dao's avatar
Tri Dao committed
189
        logits = model(input_ids).logits
Tri Dao's avatar
Tri Dao committed
190
191
        logits = rearrange(logits, "(b s) d -> b s d", b=batch_size)
        logits, _ = all_gather_raw(logits, process_group)
192
        logits = rearrange(logits, "(n b) ... d -> b ... (n d)", b=batch_size)
Tri Dao's avatar
Tri Dao committed
193
194
    del model

Tri Dao's avatar
Tri Dao committed
195
196
197
    if rank == 0:
        # Without device_map, the model is loaded on the CPU, which is very slow
        model_ref = LlamaForCausalLM.from_pretrained(
198
            Path(checkpoint_path) / f"{model_name}-hf", device_map="auto"
Tri Dao's avatar
Tri Dao committed
199
200
201
202
203
204
205
206
        )
        model_ref.eval()
        with torch.no_grad():
            out_ref = model_ref.model(input_ids).last_hidden_state.to(device=device)
            logits_ref = model_ref(input_ids).logits.to(device=device)
        del model_ref

        model_hf = LlamaForCausalLM.from_pretrained(
207
            Path(checkpoint_path) / f"{model_name}-hf", torch_dtype=dtype, device_map="auto"
Tri Dao's avatar
Tri Dao committed
208
209
210
211
212
213
214
        )
        model_hf.eval()
        with torch.no_grad():
            out_hf = model_hf.model(input_ids).last_hidden_state.to(device=device)
            logits_hf = model_hf(input_ids).logits.to(device=device)
        del model_hf

215
216
217
218
        print(f"Output max diff: {(out - out_ref).abs().max().item()}")
        print(f"Output mean diff: {(out - out_ref).abs().mean().item()}")
        print(f"HF fp16 max diff: {(out_hf - out_ref).abs().max().item()}")
        print(f"HF fp16 mean diff: {(out_hf - out_ref).abs().mean().item()}")
Tri Dao's avatar
Tri Dao committed
219
220
        assert (out - out_ref).abs().max().item() < 2 * (out_hf - out_ref).abs().max().item()

221
222
223
224
225
226
227
        print(f"Logits max diff: {(logits - logits_ref).abs().max().item()}")
        print(f"Logits mean diff: {(logits - logits_ref).abs().mean().item()}")
        print(f"HF fp16 max diff: {(logits_hf - logits_ref).abs().max().item()}")
        print(f"HF fp16 mean diff: {(logits_hf - logits_ref).abs().mean().item()}")
        assert (logits - logits_ref).abs().max().item() < 2 * (
            logits_hf - logits_ref
        ).abs().max().item()
Tri Dao's avatar
Tri Dao committed
228
229
230


# @pytest.mark.parametrize('model_name', ["7B", "13B"])
231
232
@pytest.mark.parametrize("model_name", ["7B"])
@pytest.mark.parametrize("checkpoint_format", ["meta", "hf"])
233
def test_llama_generation(model_name, checkpoint_format):
234
235
236
    checkpoint_path = (
        Path(os.environ.get("CHECKPOINT_DIR", current_dir.parent.parent / "checkpoints")) / "llama"
    )
Tri Dao's avatar
Tri Dao committed
237
238

    dtype = torch.float16
239
    device = "cuda"
240
241
    config = config_from_checkpoint(checkpoint_path, model_name, checkpoint_format)
    config = llama_config_to_gpt2_config(config)
Tri Dao's avatar
Tri Dao committed
242
243
244
245
246
247
    config.use_flash_attn = True
    config.fused_bias_fc = True
    config.fused_mlp = False  # We don't have fused GatedMLP yet
    config.fused_dropout_add_ln = True
    config.residual_in_fp32 = True

248
    tokenizer = LlamaTokenizer.from_pretrained(Path(checkpoint_path) / f"{model_name}-hf")
Tri Dao's avatar
Tri Dao committed
249
250
251
252
253
254
    eos_token_id = tokenizer.eos_token_id

    torch.manual_seed(0)
    batch_size = 1
    seqlen = 100
    max_length = 150
255
256
257
    input_ids = torch.randint(
        0, config.vocab_size, (batch_size, seqlen), dtype=torch.long, device=device
    )
Tri Dao's avatar
Tri Dao committed
258

259
260
261
    model_hf = LlamaForCausalLM.from_pretrained(
        Path(checkpoint_path) / f"{model_name}-hf", torch_dtype=dtype, device_map={"": device}
    )
Tri Dao's avatar
Tri Dao committed
262
263
264
265
    model_hf.eval()
    print("HF fp16")
    torch.cuda.synchronize()
    start = time.time()
266
267
268
    out_hf = model_hf.generate(
        input_ids=input_ids, max_length=max_length, return_dict_in_generate=True, output_scores=True
    )
Tri Dao's avatar
Tri Dao committed
269
    torch.cuda.synchronize()
270
    print(f"Prompt processing + decoding time: {(time.time() - start) * 1000:.0f}ms")
Tri Dao's avatar
Tri Dao committed
271
272
    del model_hf

Tri Dao's avatar
Tri Dao committed
273
    # Need auto here since the 13B fp32 model doesn't fit in memory on a A100 40GB
274
275
276
    model_ref = LlamaForCausalLM.from_pretrained(
        Path(checkpoint_path) / f"{model_name}-hf", device_map="auto"
    )
Tri Dao's avatar
Tri Dao committed
277
278
    model_ref.eval()
    with torch.no_grad():
279
        logits_ref = model_ref(out_hf.sequences).logits[:, (seqlen - 1) : -1].to(device=device)
Tri Dao's avatar
Tri Dao committed
280
281
    del model_ref

282
283
284
    pretrained_state_dict = _pretrained_state_dict_from_checkpoint(
        checkpoint_path, model_name, config, checkpoint_format
    )
Tri Dao's avatar
Tri Dao committed
285
    model = GPTLMHeadModel(config, device=device, dtype=dtype)
Tri Dao's avatar
Tri Dao committed
286
    model.load_state_dict(pretrained_state_dict)
Tri Dao's avatar
Tri Dao committed
287
288
    model.eval()

289
    print("Without CUDA graph")
Tri Dao's avatar
Tri Dao committed
290
291
    torch.cuda.synchronize()
    start = time.time()
292
293
294
295
296
297
298
299
300
301
    out = model.generate(
        input_ids=input_ids,
        max_length=max_length,
        eos_token_id=eos_token_id,
        fused_ft_kernel=True,
        return_dict_in_generate=True,
        output_scores=True,
        timing=True,
        teacher_outputs=out_hf.sequences,
    )
Tri Dao's avatar
Tri Dao committed
302
    torch.cuda.synchronize()
303
    print(f"Prompt processing + decoding time: {(time.time() - start) * 1000:.0f}ms")
Tri Dao's avatar
Tri Dao committed
304
305
306
307

    # Capture graph outside the timing loop
    batch_size, seqlen_og = input_ids.shape
    model._decoding_cache = update_graph_cache(model, None, batch_size, seqlen_og, max_length)
308
    print("With CUDA graph")
Tri Dao's avatar
Tri Dao committed
309
310
    torch.cuda.synchronize()
    start = time.time()
311
312
313
314
315
316
317
318
319
320
    out_cg = model.generate(
        input_ids=input_ids,
        max_length=max_length,
        fused_ft_kernel=True,
        cg=True,
        return_dict_in_generate=True,
        output_scores=True,
        timing=True,
        teacher_outputs=out_hf.sequences,
    )
Tri Dao's avatar
Tri Dao committed
321
    torch.cuda.synchronize()
322
    print(f"Prompt processing + decoding time: {(time.time() - start) * 1000:.0f}ms")
Tri Dao's avatar
Tri Dao committed
323
324

    with torch.no_grad():
325
        logits_parallel = model(out_hf.sequences).logits[:, (seqlen - 1) : -1]
Tri Dao's avatar
Tri Dao committed
326
327
328
329
330
331
332
333
    logits_hf = torch.stack(out_hf.scores, dim=1)
    logits = torch.stack(out.scores, dim=1)
    logits_cg = torch.stack(out_cg.scores, dim=1)

    del model

    hf_error = (logits_hf - logits_ref).abs().max().item()

334
335
336
    print(f"HF fp16 logits max diff: {hf_error}")
    print(f"Logits max diff: {(logits - logits_ref).abs().max().item()}")
    print(f"Logits CG max diff: {(logits_cg - logits_ref).abs().max().item()}")
337
338
339

    assert (logits_parallel - logits_ref).abs().max().item() < 2 * hf_error
    assert (logits - logits_ref).abs().max().item() < 2 * hf_error
Tri Dao's avatar
Tri Dao committed
340
    assert torch.equal(logits_cg, logits)
Tri Dao's avatar
Tri Dao committed
341
342
343


# torchrun --no_python --nproc_per_node=2 pytest -q -s tests/models/test_llama.py -k "llama_parallel_generation"
344
345
346
@pytest.mark.parametrize("world_size", [2])
@pytest.mark.parametrize("model_name", ["13B"])
@pytest.mark.parametrize("checkpoint_format", ["meta", "hf"])
347
def test_llama_parallel_generation(model_name, world_size, checkpoint_format):
Tri Dao's avatar
Tri Dao committed
348
349
350
351
352
353
    """Check that our implementation matches the HF implementation:
    the scores in fp16 should be around the same as the HF scores in fp16, when compared to
    the HF scores in fp32.
    """
    from apex.transformer import parallel_state

354
355
356
    checkpoint_path = (
        Path(os.environ.get("CHECKPOINT_DIR", current_dir.parent.parent / "checkpoints")) / "llama"
    )
Tri Dao's avatar
Tri Dao committed
357
358

    dtype = torch.float16
359
360
    config = config_from_checkpoint(checkpoint_path, model_name, checkpoint_format)
    config = llama_config_to_gpt2_config(config)
Tri Dao's avatar
Tri Dao committed
361
362
363
364
365
366
367
368
369
370
    config.use_flash_attn = False
    config.fused_bias_fc = True
    config.fused_mlp = False  # We don't have fused GatedMLP yet
    config.fused_dropout_add_ln = False
    config.residual_in_fp32 = True
    config.pad_vocab_size_multiple = 8 * world_size
    config.sequence_parallel = False  # Need to set this to False for generation

    os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "0"
    if not torch.distributed.is_initialized():
371
372
        torch.distributed.init_process_group(backend="nccl", init_method="env://")
    device = f"cuda:{torch.distributed.get_rank()}"
Tri Dao's avatar
Tri Dao committed
373
374
375
376
377
378
379
380
381
    assert world_size <= torch.distributed.get_world_size()
    parallel_state.initialize_model_parallel(tensor_model_parallel_size_=world_size)
    rank = parallel_state.get_tensor_model_parallel_rank()
    process_group = parallel_state.get_tensor_model_parallel_group()

    torch.manual_seed(0)
    batch_size = 1
    seqlen = 100
    max_length = 150
382
383
384
    input_ids = torch.randint(
        0, config.vocab_size, (batch_size, seqlen), dtype=torch.long, device=device
    )
Tri Dao's avatar
Tri Dao committed
385
386
387
388
389

    # Need this, otherwise when we capture the graph the process for GPU 1 would run on both
    # GPU0 and GPU1 and things would hang
    torch.cuda.set_device(device)

390
391
392
    pretrained_state_dict = _pretrained_state_dict_from_checkpoint(
        checkpoint_path, model_name, config, checkpoint_format
    )
Tri Dao's avatar
Tri Dao committed
393
394
395
396
    model = GPTLMHeadModel(config, process_group=process_group, device=device, dtype=dtype)
    model.load_state_dict(shard_state_dict_tp(pretrained_state_dict, config, world_size, rank))
    model.eval()

397
    print("Without CUDA graph")
Tri Dao's avatar
Tri Dao committed
398
    out = model.generate(
399
400
401
402
403
        input_ids=input_ids,
        max_length=max_length,
        tensor_parallel=world_size,
        vocab_size=config.vocab_size,
        fused_ft_kernel=True,
Tri Dao's avatar
Tri Dao committed
404
        # teacher_outputs=out_hf.sequences,
405
406
407
        return_dict_in_generate=True,
        output_scores=True,
        timing=True,
Tri Dao's avatar
Tri Dao committed
408
409
410
411
412
    )

    # Capture graph outside the timing loop
    batch_size, seqlen_og = input_ids.shape
    model._decoding_cache = update_graph_cache(model, None, batch_size, seqlen_og, max_length)
413
    print("With CUDA graph")
Tri Dao's avatar
Tri Dao committed
414
    out_cg = model.generate(
415
416
417
418
419
420
        input_ids=input_ids,
        max_length=max_length,
        tensor_parallel=world_size,
        vocab_size=config.vocab_size,
        fused_ft_kernel=True,
        cg=True,
Tri Dao's avatar
Tri Dao committed
421
        # teacher_outputs=out_hf.sequences,
422
423
424
        return_dict_in_generate=True,
        output_scores=True,
        timing=True,
Tri Dao's avatar
Tri Dao committed
425
426
427
428
429
430
431
    )
    del model
    parallel_state.destroy_model_parallel()

    if rank == 0:
        # Without device_map, the model is loaded on the CPU, which is very slow
        model_hf = LlamaForCausalLM.from_pretrained(
432
            Path(checkpoint_path) / f"{model_name}-hf", torch_dtype=dtype, device_map="auto"
Tri Dao's avatar
Tri Dao committed
433
434
435
436
437
438
439
        )
        model_hf.eval()
        print("HF fp16")
        torch.cuda.synchronize()
        start = time.time()
        with torch.inference_mode():
            out_hf = model_hf.generate(
440
441
442
443
                input_ids=input_ids,
                max_length=max_length,
                return_dict_in_generate=True,
                output_scores=True,
Tri Dao's avatar
Tri Dao committed
444
445
            )
        torch.cuda.synchronize()
446
        print(f"Prompt processing + decoding time: {(time.time() - start) * 1000:.0f}ms")
Tri Dao's avatar
Tri Dao committed
447
448
449
        del model_hf

        model_ref = LlamaForCausalLM.from_pretrained(
450
            Path(checkpoint_path) / f"{model_name}-hf", device_map="auto"
Tri Dao's avatar
Tri Dao committed
451
452
453
        )
        model_ref.eval()
        with torch.inference_mode():
454
            logits_ref = model_ref(out_hf.sequences).logits[:, (seqlen - 1) : -1]
Tri Dao's avatar
Tri Dao committed
455
456
457
458
459
460
461
        del model_ref
        logits_hf = torch.stack(out_hf.scores, dim=1)

        logits = torch.stack(out.scores, dim=1)
        logits_cg = torch.stack(out_cg.scores, dim=1)

        hf_error = (logits_hf - logits_ref).abs().max().item()
462
463
        print(f"HF fp16 logits max diff: {hf_error}")
        print(f"Logits max diff: {(logits - logits_ref).abs().max().item()}")
Tri Dao's avatar
Tri Dao committed
464
        assert (logits - logits_ref).abs().max().item() < 2 * hf_error
465
        print(f"Logits CG max diff: {(logits_cg - logits_ref).abs().max().item()}")
Tri Dao's avatar
Tri Dao committed
466
        assert torch.equal(logits_cg, logits)
467
468
469


@torch.no_grad()
470
@pytest.mark.parametrize("world_size", [2])
471
472
473
def test_llama_parallel_uneven_num_heads(world_size):
    from apex.transformer import parallel_state

474
475
476
    checkpoint_path = (
        Path(os.environ.get("CHECKPOINT_DIR", current_dir.parent.parent / "checkpoints")) / "llama"
    )
477
    num_attention_heads = world_size + 1
478
    model_name = f"teeny-{num_attention_heads}-heads"
479
480

    if not torch.distributed.is_initialized():
481
482
        torch.distributed.init_process_group(backend="nccl", init_method="env://")
    device = f"cuda:{torch.distributed.get_rank()}"
483
484
485
486
487
488
489
    assert world_size <= torch.distributed.get_world_size()
    parallel_state.initialize_model_parallel(tensor_model_parallel_size_=world_size)
    rank = parallel_state.get_tensor_model_parallel_rank()
    process_group = parallel_state.get_tensor_model_parallel_group()

    dtype = torch.float16
    llama_config = LlamaConfig(
490
491
        hidden_size=256
        * num_attention_heads,  # ParallelGatedMlp hidden_features must be divisible by 256
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
        intermediate_size=256 * num_attention_heads * 4,
        num_hidden_layers=4,
        num_attention_heads=num_attention_heads,
        initializer_range=0.5,  # Set crazy init range so we don't have near zero weights implying a vacuous test.
    )
    config = llama_config_to_gpt2_config(llama_config)
    config.use_flash_attn = True
    config.fused_bias_fc = True
    config.fused_mlp = False  # We don't have fused GatedMLP yet
    config.fused_dropout_add_ln = True
    config.residual_in_fp32 = True

    torch.manual_seed(0)
    batch_size = 2
    max_seqlen = 256
    seqlens = torch.randint(max_seqlen // 2, max_seqlen + 1, (batch_size,), device=device)
508
509
510
    input_ids = torch.randint(
        0, config.vocab_size, (batch_size, max_seqlen), dtype=torch.long, device=device
    )
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531

    # Create a shared test model.
    if rank == 0:
        LlamaForCausalLM(config=llama_config).save_pretrained(checkpoint_path / f"{model_name}-hf")
    torch.distributed.barrier()

    # Run the standard forward pass test.
    pretrained_state_dict = _pretrained_state_dict_from_checkpoint(
        checkpoint_path, model_name, config, checkpoint_format="hf"
    )
    model = GPTLMHeadModel(config, process_group=process_group, device=device, dtype=dtype)
    model.load_state_dict(shard_state_dict_tp(pretrained_state_dict, config, world_size, rank))
    model.eval()

    # TODO: Avoid duplicate code. Modularize the comparison of two forward pass diffs.
    out = model.transformer(input_ids)
    out, _ = all_gather_raw(out, process_group=process_group)
    out = rearrange(out, "(b s) d -> b s d", b=batch_size)
    logits = model(input_ids).logits
    logits = rearrange(logits, "(b s) d -> b s d", b=batch_size)
    logits, _ = all_gather_raw(logits, process_group)
532
    logits = rearrange(logits, "(n b) ... d -> b ... (n d)", b=batch_size)
533
534
535

    if rank == 0:
        model_ref = LlamaForCausalLM.from_pretrained(
536
            Path(checkpoint_path) / f"{model_name}-hf", device_map="auto"
537
538
539
540
541
542
543
        )
        model_ref.eval()
        out_ref = model_ref.model(input_ids).last_hidden_state.to(device=device)
        logits_ref = model_ref(input_ids).logits.to(device=device)
        del model_ref

        model_hf = LlamaForCausalLM.from_pretrained(
544
            Path(checkpoint_path) / f"{model_name}-hf", torch_dtype=dtype, device_map="auto"
545
546
547
548
549
550
        )
        model_hf.eval()
        out_hf = model_hf.model(input_ids).last_hidden_state.to(device=device)
        logits_hf = model_hf(input_ids).logits.to(device=device)
        del model_hf

551
552
553
554
        print(f"Output max diff: {(out - out_ref).abs().max().item()}")
        print(f"Output mean diff: {(out - out_ref).abs().mean().item()}")
        print(f"HF fp16 max diff: {(out_hf - out_ref).abs().max().item()}")
        print(f"HF fp16 mean diff: {(out_hf - out_ref).abs().mean().item()}")
555
556
        assert (out - out_ref).abs().max().item() < 2 * (out_hf - out_ref).abs().max().item()

557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
        print(f"Logits max diff: {(logits - logits_ref).abs().max().item()}")
        print(f"Logits mean diff: {(logits - logits_ref).abs().mean().item()}")
        print(f"HF fp16 max diff: {(logits_hf - logits_ref).abs().max().item()}")
        print(f"HF fp16 mean diff: {(logits_hf - logits_ref).abs().mean().item()}")
        assert (logits - logits_ref).abs().max().item() < 2 * (
            logits_hf - logits_ref
        ).abs().max().item()

        if os.path.exists(checkpoint_path / f"{model_name}-hf"):
            shutil.rmtree(checkpoint_path / f"{model_name}-hf")


@torch.no_grad()
def test_inv_remap_state_dict_hf_llama():
    checkpoint_path = (
        Path(os.environ.get("CHECKPOINT_DIR", current_dir.parent.parent / "checkpoints")) / "llama"
    )
    model_name = f"teeny"

    llama_config = LlamaConfig(
        num_attention_heads=2,
        hidden_size=256 * 2,
        intermediate_size=256 * 2 * 4,
        num_hidden_layers=4,
    )
    config = llama_config_to_gpt2_config(llama_config)
    config.use_flash_attn = True
    config.fused_bias_fc = True
    config.fused_mlp = False  # We don't have fused GatedMLP yet
    config.fused_dropout_add_ln = True
    config.residual_in_fp32 = True

    # Set up.
    LlamaForCausalLM(config=llama_config).save_pretrained(checkpoint_path / f"{model_name}-hf")

    # inv_remap_state_dict_hf_llama should be the inverse of remap_state_dict_hf_llama
    state_dict = state_dict_from_pretrained(checkpoint_path / f"{model_name}-hf")
    state_dict = {key: val for key, val in state_dict.items() if "rotary_emb.inv_freq" not in key}
    pretrained_state_dict = remap_state_dict_hf_llama(state_dict, config)
    state_dict_recover = inv_remap_state_dict_hf_llama(pretrained_state_dict, config)

    assert set(state_dict_recover.keys()) == set(state_dict.keys())

    for key in state_dict_recover.keys():
        torch.testing.assert_close(state_dict_recover[key], state_dict[key])
602

603
604
605
    # Tear down.
    if os.path.exists(checkpoint_path / f"{model_name}-hf"):
        shutil.rmtree(checkpoint_path / f"{model_name}-hf")