403 lines
14 KiB
Python
403 lines
14 KiB
Python
import argparse
|
|
from dataclasses import dataclass
|
|
import json
|
|
import os
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
from torch.utils.data import DataLoader, Subset
|
|
from transformers import AutoModel, AutoTokenizer
|
|
|
|
from deepspec.data import ConversationCollator
|
|
from deepspec.data.target_cache_dataset import (
|
|
AsyncTargetCacheWriter,
|
|
LocalCacheWriteSummary,
|
|
atomic_json_dump,
|
|
build_global_target_cache_shard_map,
|
|
build_target_cache_manifest,
|
|
cleanup_target_cache_tmp_dir,
|
|
compute_local_sample_range,
|
|
finalize_target_cache_index,
|
|
load_local_cache_write_summary,
|
|
prepare_target_cache_output_dir,
|
|
rename_local_target_cache_shards,
|
|
write_target_cache_manifest,
|
|
)
|
|
from deepspec.data.jsonl_dataset import JsonLineDataset
|
|
from deepspec.utils import (
|
|
CustomJSONEncoder,
|
|
get_git_diff,
|
|
get_git_sha,
|
|
init_dist,
|
|
is_global_main_process,
|
|
load_config,
|
|
main_process_first,
|
|
parse_opts_to_config,
|
|
print_on_global_main,
|
|
print_on_local_main,
|
|
seed_all,
|
|
)
|
|
|
|
os.environ["USE_TORCH"] = "true"
|
|
os.environ["WANDB_DISABLED"] = "true"
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
|
|
# PyTorch 2.10 Inductor still reads the legacy allow_tf32 flag while compiling.
|
|
torch.set_float32_matmul_precision("high")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TargetForwardResult:
|
|
target_hidden_states: torch.Tensor
|
|
target_last_hidden_states: torch.Tensor
|
|
|
|
|
|
def _get_target_backbone(target_model):
|
|
model_type = str(target_model.config.model_type)
|
|
if model_type in ("gemma4", "gemma4_unified"):
|
|
if hasattr(target_model, "language_model"):
|
|
return target_model.language_model
|
|
if hasattr(target_model, "model") and hasattr(target_model.model, "language_model"):
|
|
return target_model.model.language_model
|
|
assert False, "Gemma4 target model must expose a text language_model."
|
|
return getattr(target_model, "model", target_model)
|
|
|
|
|
|
def _get_target_hidden_size(target_model) -> int:
|
|
model_type = str(target_model.config.model_type)
|
|
if model_type in ("gemma4", "gemma4_unified"):
|
|
return int(target_model.config.text_config.hidden_size)
|
|
return int(target_model.config.hidden_size)
|
|
|
|
|
|
def _get_hook_tensor(output):
|
|
if isinstance(output, torch.Tensor):
|
|
return output
|
|
if isinstance(output, (tuple, list)) and output:
|
|
first = output[0]
|
|
if isinstance(first, torch.Tensor):
|
|
return first
|
|
raise TypeError(f"Unsupported target hook output type: {type(output)!r}")
|
|
|
|
|
|
def run_target_forward_with_hooks(
|
|
*,
|
|
target_model,
|
|
input_ids: torch.Tensor,
|
|
attention_mask: torch.Tensor,
|
|
target_layer_ids,
|
|
):
|
|
backbone = _get_target_backbone(target_model)
|
|
layer_modules = backbone.layers
|
|
target_layer_ids = [int(layer_id) for layer_id in target_layer_ids]
|
|
captured_hidden_states = {}
|
|
handles = []
|
|
|
|
def capture_layer(layer_id: int):
|
|
def hook(_module, _inputs, output):
|
|
captured_hidden_states[layer_id] = _get_hook_tensor(output).detach()
|
|
|
|
return hook
|
|
|
|
try:
|
|
if -1 in target_layer_ids:
|
|
handles.append(
|
|
backbone.embed_tokens.register_forward_hook(capture_layer(-1))
|
|
)
|
|
for layer_id in target_layer_ids:
|
|
if layer_id < 0:
|
|
continue
|
|
handles.append(
|
|
layer_modules[layer_id].register_forward_hook(capture_layer(layer_id))
|
|
)
|
|
|
|
with torch.no_grad():
|
|
target_output = target_model(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
output_hidden_states=False,
|
|
use_cache=False,
|
|
)
|
|
target_last_hidden_states = target_output.last_hidden_state.detach()
|
|
target_hidden_states = torch.cat(
|
|
[captured_hidden_states[layer_id] for layer_id in target_layer_ids],
|
|
dim=-1,
|
|
)
|
|
finally:
|
|
for handle in handles:
|
|
handle.remove()
|
|
captured_hidden_states.clear()
|
|
|
|
return TargetForwardResult(
|
|
target_hidden_states=target_hidden_states,
|
|
target_last_hidden_states=target_last_hidden_states,
|
|
)
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", required=True)
|
|
parser.add_argument("--opts", action="append", default=[])
|
|
parser.add_argument(
|
|
"--train-data-path",
|
|
action="append",
|
|
required=True,
|
|
help="Training JSONL path. Repeat this argument to use multiple files.",
|
|
)
|
|
parser.add_argument("--output-dir", required=True)
|
|
parser.add_argument("--min-loss-tokens", type=int, default=14)
|
|
parser.add_argument("--max-shard-bytes", type=int, default=64 * 1024**3)
|
|
parser.add_argument("--local-batch-size", type=int, default=32)
|
|
parser.add_argument("--num-workers", type=int, default=4)
|
|
cli_args = parser.parse_args()
|
|
config = parse_opts_to_config(cli_args.opts, load_config(cli_args.config))
|
|
return cli_args, config
|
|
|
|
|
|
def _write_manifest(
|
|
*,
|
|
output_dir: str,
|
|
config,
|
|
train_data_paths,
|
|
target_layer_ids,
|
|
hidden_size: int,
|
|
min_loss_tokens: int,
|
|
shards,
|
|
):
|
|
num_samples = sum(
|
|
int(
|
|
load_local_cache_write_summary(
|
|
os.path.join(output_dir, "_tmp", f"rank_{rank}")
|
|
)["num_local_samples"]
|
|
)
|
|
for rank in range(dist.get_world_size())
|
|
)
|
|
manifest = build_target_cache_manifest(
|
|
num_samples=num_samples,
|
|
shards=shards,
|
|
target_layer_ids=target_layer_ids,
|
|
hidden_size=hidden_size,
|
|
extra_fields={
|
|
"target_model_name_or_path": str(config.model.target_model_name_or_path),
|
|
"source_jsonl_paths": [str(path) for path in train_data_paths],
|
|
"chat_template": str(config.data.chat_template),
|
|
"max_length": int(config.data.max_length),
|
|
"min_loss_tokens": int(min_loss_tokens),
|
|
"project_name": (
|
|
str(config.get("project_name"))
|
|
if config.get("project_name") is not None
|
|
else None
|
|
),
|
|
"exp_name": (
|
|
str(config.get("exp_name"))
|
|
if config.get("exp_name") is not None
|
|
else None
|
|
),
|
|
"git_sha": str(get_git_sha()),
|
|
},
|
|
)
|
|
write_target_cache_manifest(output_dir=output_dir, manifest=manifest)
|
|
|
|
|
|
def _print_prepare_progress(*, global_rank: int, processed_samples: int, total_samples: int):
|
|
print(
|
|
f"[prepare rank {global_rank}] {processed_samples}/{total_samples} samples",
|
|
flush=True,
|
|
)
|
|
|
|
|
|
def main(local_rank: int):
|
|
cli_args, config = parse_args()
|
|
train_data_paths = list(cli_args.train_data_path)
|
|
target_layer_ids = [int(layer_id) for layer_id in config.model.target_layer_ids]
|
|
min_loss_tokens = int(cli_args.min_loss_tokens)
|
|
seed_all(int(config.seed))
|
|
device, global_rank, world_size = init_dist(local_rank)
|
|
output_dir = os.path.abspath(cli_args.output_dir)
|
|
print_on_local_main(json.dumps(config, indent=4, cls=CustomJSONEncoder), flush=True)
|
|
print_on_local_main(
|
|
json.dumps(
|
|
{
|
|
"train_data_path": train_data_paths,
|
|
"output_dir": output_dir,
|
|
"target_layer_ids": target_layer_ids,
|
|
"min_loss_tokens": min_loss_tokens,
|
|
"max_shard_bytes": int(cli_args.max_shard_bytes),
|
|
"local_batch_size": int(cli_args.local_batch_size),
|
|
"num_workers": int(cli_args.num_workers),
|
|
},
|
|
indent=4,
|
|
),
|
|
flush=True,
|
|
)
|
|
if global_rank == 0:
|
|
prepare_target_cache_output_dir(output_dir)
|
|
dist.barrier()
|
|
|
|
rank_dir = os.path.join(output_dir, "_tmp", f"rank_{global_rank}")
|
|
os.makedirs(rank_dir, exist_ok=True)
|
|
|
|
with main_process_first():
|
|
dataset = JsonLineDataset(data_paths=train_data_paths)
|
|
|
|
local_start, local_end = compute_local_sample_range(
|
|
num_samples=len(dataset),
|
|
rank=global_rank,
|
|
world_size=world_size,
|
|
)
|
|
local_total_samples = local_end - local_start
|
|
|
|
local_subset = Subset(dataset, range(local_start, local_end))
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
config.model.target_model_name_or_path,
|
|
)
|
|
target_model = AutoModel.from_pretrained(
|
|
config.model.target_model_name_or_path,
|
|
dtype=torch.bfloat16,
|
|
attn_implementation="sdpa",
|
|
).to(device=device).eval()
|
|
target_hidden_size = _get_target_hidden_size(target_model)
|
|
train_collator = ConversationCollator(
|
|
tokenizer=tokenizer,
|
|
chat_template=config.data.chat_template,
|
|
max_length=config.data.max_length,
|
|
min_loss_tokens=min_loss_tokens,
|
|
)
|
|
dataloader = DataLoader(
|
|
local_subset,
|
|
batch_size=int(cli_args.local_batch_size),
|
|
collate_fn=train_collator,
|
|
num_workers=int(cli_args.num_workers),
|
|
pin_memory=True,
|
|
drop_last=False,
|
|
)
|
|
writer = AsyncTargetCacheWriter(
|
|
rank_dir=rank_dir,
|
|
max_shard_bytes=int(cli_args.max_shard_bytes),
|
|
max_queue_size=int(cli_args.local_batch_size) * 4,
|
|
)
|
|
|
|
processed_local_samples = 0
|
|
last_progress_printed = 0
|
|
try:
|
|
with torch.no_grad():
|
|
for batch_idx, batch in enumerate(dataloader):
|
|
processed_local_samples = min(
|
|
(batch_idx + 1) * int(cli_args.local_batch_size),
|
|
local_total_samples,
|
|
)
|
|
should_print_progress = (
|
|
processed_local_samples - last_progress_printed >= 100
|
|
or processed_local_samples == local_total_samples
|
|
)
|
|
if batch is None:
|
|
if should_print_progress:
|
|
_print_prepare_progress(
|
|
global_rank=global_rank,
|
|
processed_samples=processed_local_samples,
|
|
total_samples=local_total_samples,
|
|
)
|
|
last_progress_printed = processed_local_samples
|
|
continue
|
|
batch = {
|
|
key: value.to(device, non_blocking=True)
|
|
for key, value in batch.items()
|
|
}
|
|
target_result = run_target_forward_with_hooks(
|
|
target_model=target_model,
|
|
input_ids=batch["input_ids"],
|
|
attention_mask=batch["attention_mask"],
|
|
target_layer_ids=target_layer_ids,
|
|
)
|
|
seq_lens = batch["attention_mask"].sum(dim=1).tolist()
|
|
for sample_idx_in_batch, seq_len in enumerate(seq_lens):
|
|
seq_len = int(seq_len)
|
|
writer.write_sample(
|
|
input_ids=batch["input_ids"][sample_idx_in_batch, :seq_len],
|
|
attention_mask=batch["attention_mask"][
|
|
sample_idx_in_batch, :seq_len
|
|
],
|
|
loss_mask=batch["loss_mask"][sample_idx_in_batch, :seq_len],
|
|
target_hidden_states=target_result.target_hidden_states[
|
|
sample_idx_in_batch, :seq_len
|
|
],
|
|
target_last_hidden_states=target_result.target_last_hidden_states[
|
|
sample_idx_in_batch, :seq_len
|
|
],
|
|
)
|
|
if should_print_progress:
|
|
_print_prepare_progress(
|
|
global_rank=global_rank,
|
|
processed_samples=processed_local_samples,
|
|
total_samples=local_total_samples,
|
|
)
|
|
last_progress_printed = processed_local_samples
|
|
finally:
|
|
writer.close()
|
|
del target_model
|
|
torch.cuda.empty_cache()
|
|
dataset.close()
|
|
summary = LocalCacheWriteSummary(
|
|
global_rank=global_rank,
|
|
source_sample_start=local_start,
|
|
source_sample_end=local_end,
|
|
num_local_samples=writer.num_local_samples,
|
|
num_local_shards=len(writer.local_shard_files),
|
|
local_shard_files=list(writer.local_shard_files),
|
|
)
|
|
atomic_json_dump(summary.to_json(), os.path.join(rank_dir, "summary.json"))
|
|
dist.barrier()
|
|
|
|
shard_map = None
|
|
summaries = None
|
|
if is_global_main_process():
|
|
summaries = [
|
|
load_local_cache_write_summary(
|
|
os.path.join(output_dir, "_tmp", f"rank_{rank}")
|
|
)
|
|
for rank in range(world_size)
|
|
]
|
|
shard_map, shards = build_global_target_cache_shard_map(summaries)
|
|
broadcast_payload = [shard_map]
|
|
dist.broadcast_object_list(broadcast_payload, src=0)
|
|
shard_map = broadcast_payload[0]
|
|
local_summary = load_local_cache_write_summary(rank_dir)
|
|
rename_local_target_cache_shards(
|
|
output_dir=output_dir,
|
|
rank_dir=rank_dir,
|
|
summary=local_summary,
|
|
shard_map=shard_map,
|
|
)
|
|
dist.barrier()
|
|
|
|
if is_global_main_process():
|
|
assert summaries is not None
|
|
num_valid_samples = finalize_target_cache_index(
|
|
output_dir=output_dir,
|
|
summaries=summaries,
|
|
shard_map=shard_map,
|
|
)
|
|
_write_manifest(
|
|
output_dir=output_dir,
|
|
config=config,
|
|
train_data_paths=train_data_paths,
|
|
target_layer_ids=target_layer_ids,
|
|
hidden_size=target_hidden_size,
|
|
min_loss_tokens=min_loss_tokens,
|
|
shards=shards,
|
|
)
|
|
cleanup_target_cache_tmp_dir(output_dir)
|
|
print_on_global_main(
|
|
f"Prepared target cache at {output_dir} with "
|
|
f"{num_valid_samples}/{len(dataset)} valid samples."
|
|
)
|
|
dist.barrier()
|
|
dist.destroy_process_group()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if os.path.exists(".git"):
|
|
print(f"git status:", "\n\n".join(get_git_sha(detail_info=True)))
|
|
print("git diff:", get_git_diff())
|
|
torch.multiprocessing.spawn(main, nprocs=torch.cuda.device_count())
|