CS336 Lab 1 实验笔记

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44

class PreTokenizer(ABC):
    @staticmethod
    def _merge_pre_token_counts(*pre_token_counts: PreTokenCount) -> PreTokenCount:
        """Merge multiple PreTokenCount dictionaries into one.

        Returns:
            PreTokenCount: The merged PreTokenCount.
        """

    def _process_chunk(self, chunk: Chunk, special_tokens: list[Token]) -> PreTokenCount:
        """
        Process a single chunk of text and return the pre-token counts.

        Args:
            chunk (Chunk): The chunk of text to process.

        Returns:
            PreTokenCount: A dictionary-like object mapping pre-tokens to their counts.
        """

    def pre_tokenize(self, str_bytes: bytes, special_token_list: list[Token]) -> Iterator[Token]:
        """
        Pre-tokenize the given bytes string.

        Args:
            str_bytes (bytes): The input bytes string to pre-tokenize.
            special_token_list (list[Token]): The list of special tokens.
        Returns:
            Iterator[Token]: An iterator over the pre-tokens.
        """

    @abstractmethod
    def __call__(self, corpos_path: str, split_special_token: Token, special_tokens: list[Token]) -> PreTokenCount:
        """
        Pre-tokenize the given corpus.

        Args:
            corpos_path (str): Path to the corpus file.
            split_special_token (token): The special token used to split the corpus.
            special_tokens (list[Token]): List of special tokens.

        Returns:
            PreTokenCount: A dictionary-like object mapping pre-tokens to their counts.
        """

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

class MultiProcessPreTokenizer(PreTokenizer):
    def _process_chunk_with_boundry(
        self, corpos_path: str, start: int, end: int, special_tokens: list[Token]
    ) -> PreTokenCount:
        with open(corpos_path, mode="br") as f:
            f.seek(start)
            chunk = f.read(end - start)
            pre_token_count = self._process_chunk(chunk, special_tokens)
        return pre_token_count

    def __call__(self, corpos_path: str, split_special_token: Token, special_tokens: list[Token]) -> PreTokenCount:
        final_pre_token_count: PreTokenCount = defaultdict(int)

        start_time = time.time()
        file_size = os.path.getsize(corpos_path)
        num_cpus = cpu_count()

        desired_chunks = num_cpus * 100

        chunk_boundaries = find_chunk_boundaries(
            file_path=corpos_path,
            desired_num_chunks=desired_chunks,
            split_special_token=split_special_token,
        )

        chunks_args = []
        for i in range(len(chunk_boundaries) - 1):
            start = chunk_boundaries[i]
            end = chunk_boundaries[i + 1]
            chunks_args.append((corpos_path, start, end, special_tokens))

        logger.info(f"Splitting task into {len(chunks_args)} chunks.")

        with Pool(processes=num_cpus * 2) as pool:
            chunk_iter = pool.imap_unordered(self._worker_wrapper, chunks_args)

            for chunk_result in tqdm.tqdm(chunk_iter, total=len(chunks_args), desc="Pre-tokenizing"):
                for token, count in chunk_result.items():
                    final_pre_token_count[token] += count

        end_time = time.time()
        logger.info(
            "Takes {:.2f} seconds to pre-tokenize, speed: {:.2f} bytes/second",
            end_time - start_time,
            file_size / (end_time - start_time),
        )
        return final_pre_token_count

    @staticmethod
    def _worker_wrapper(args):
        tokenizer_instance = MultiProcessPreTokenizer()
        return tokenizer_instance._process_chunk_with_boundry(*args)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

def train(self) -> tuple[Vocab, list[tuple[bytes, bytes]]]:
	self._init()

	num_merges_needed = self.target_vocab_size - self.vocab_size
	if num_merges_needed <= 0:
		return {}, self.merges

	while self.vocab_size < self.target_vocab_size:
		merge_pair = self._determine_merge_pair()
		if merge_pair is None:
			break

		self.merges.append(merge_pair)
		self._merge_pair(merge_pair)

	return self.vocab, self.merges

1
2
3
4
5
6
7
8

def load_batch(dataset: npt.NDArray, batch_size: int, context_size: int, device: torch.device):
    max_start_index = dataset.shape[0] - context_size + 1
    start_indices = np.random.randint(0, max_start_index - 1, size=batch_size)
    x_batch = np.stack([dataset[i : i + context_size] for i in start_indices])
    y_batch = np.stack([dataset[i + 1 : i + context_size + 1] for i in start_indices])
    x_tensor = torch.tensor(x_batch, device=device, dtype=torch.int)
    y_tensor = torch.tensor(y_batch, device=device, dtype=torch.int)
    return x_tensor, y_tensor

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

def save_checkpoint(
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    iteration: int,
    out: str | os.PathLike | typing.BinaryIO | typing.IO[bytes],
):
    states = {
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "iteration": iteration,
    }
    torch.save(states, out)


def load_checkpoint(
    checkpoint: str | os.PathLike | typing.BinaryIO | typing.IO[bytes],
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer | None = None,
) -> int:
    states = torch.load(checkpoint)
    model.load_state_dict(states["model_state_dict"])
    if optimizer is not None:
        optimizer.load_state_dict(states["optimizer_state_dict"])
    iteration = states["iteration"]
    return iteration

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53

model:
  vocab_size: 10000
  # 上下文长度/序列最大长度 (Context length)
  context_length: 256
  # Transformer 块层数 (Number of Transformer layers)
  num_layers: 4
  # 隐藏层维度/模型大小 (Hidden dimension)
  hidden_dim: 512
  # Transformer 内部前馈网络的维度 (Inner dimension of FFN, typically 4*hidden_dim)
  inner_dim: 1344
  # 多头注意力机制的头数 (Number of attention heads)
  num_heads: 16
  # RoPE (Rotary Position Embedding) 的旋转基数（通常为 10000）
  theta: 10000.0
  # 训练设备 (Device for training: 'cuda', 'mps', or 'cpu')
  device: "cuda"
  # 模型参数的数据类型 (Data type: 'float32' or 'bfloat16')
  dtype: "bfloat16"

optimizer:
  # 学习率 (Learning rate)
  learning_rate: 3.0e-4
  # 权重衰减 (Weight decay for AdamW)
  weight_decay: 1e-2
  # AdamW 优化器的一阶矩估计衰减率 (Beta1)
  beta1: 0.9
  # AdamW 优化器的二阶矩估计衰减率 (Beta2)
  beta2: 0.999
  # AdamW 优化器的数值稳定性参数 (Epsilon)
  eps: 1.0e-8

training:
  # 批次大小 (Batch size)
  batch_size: 96
  # 总训练迭代次数 (Total training iterations)
  total_iterations: 20000
  # 学习率预热的迭代次数 (Warmup iterations)
  warmup_iterations: 1000
  # 学习率余弦退火周期迭代次数 (Cosine cycle iterations)
  cosine_cycle_iterations: 20000
  # 梯度裁剪的最大 L2 范数 (Max L2 norm for gradient clipping)
  max_l2_norm: 1.0
  # 模型保存的间隔（迭代次数） (Checkpoint saving interval)
  checkpoint_interval: 1000
  # 模型和配置的输出目录 (Output directory)
  output_dir: "./output_full/"
  # 训练数据文件路径 (Path to training data .npy file)
  train_data: "save/data/TS-train.bin"
  # 验证数据文件路径 (Path to validation data .npy file)
  val_data: "save/data/owt_train.bin"
  # 保存检查点的步长（在 `train` 函数中用于判断是否保存检查点）
  save_step: 500
  log_step: 50