|
|
import os |
|
|
import json |
|
|
import logging |
|
|
import threading |
|
|
import uuid |
|
|
import time |
|
|
import sys |
|
|
import gc |
|
|
import multiprocessing |
|
|
import shutil |
|
|
import math |
|
|
from datetime import datetime |
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed |
|
|
from itertools import chain |
|
|
|
|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.nn.functional as F |
|
|
import gradio as gr |
|
|
import transformers |
|
|
import datasets |
|
|
from dotenv import load_dotenv |
|
|
from datasets import load_dataset, get_dataset_config_names, IterableDataset |
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, TrainerCallback, AutoConfig, DataCollatorForLanguageModeling |
|
|
from huggingface_hub import login, whoami, create_repo, upload_folder |
|
|
import spaces |
|
|
|
|
|
try: |
|
|
load_dotenv() |
|
|
except: |
|
|
pass |
|
|
|
|
|
transformers.logging.set_verbosity_error() |
|
|
datasets.logging.set_verbosity_error() |
|
|
logging.getLogger("transformers").setLevel(logging.CRITICAL) |
|
|
logging.getLogger("datasets").setLevel(logging.CRITICAL) |
|
|
logging.getLogger("torch").setLevel(logging.CRITICAL) |
|
|
logging.basicConfig(level=logging.CRITICAL, stream=sys.stderr) |
|
|
|
|
|
if torch.cuda.is_available(): |
|
|
torch.backends.cuda.matmul.allow_tf32 = True |
|
|
torch.backends.cudnn.allow_tf32 = True |
|
|
torch.backends.cudnn.benchmark = True |
|
|
|
|
|
JOBS = {} |
|
|
|
|
|
def activation_quant(x): |
|
|
scale = 127.0 / x.abs().max(dim=-1, keepdim=True).values.clamp_(min=1e-5) |
|
|
y = (x * scale).round().clamp_(-128, 127) / scale |
|
|
return y + x - x.detach() |
|
|
|
|
|
def weight_quant(w): |
|
|
scale = 1.0 / w.abs().mean().clamp_(min=1e-5) |
|
|
u = (w * scale).round().clamp_(-1, 1) / scale |
|
|
return u + w - w.detach() |
|
|
|
|
|
class BitLinear(nn.Linear): |
|
|
def forward(self, x): |
|
|
target_dtype = x.dtype |
|
|
|
|
|
w = self.weight.to(target_dtype) |
|
|
w_quant = weight_quant(w).to(target_dtype) |
|
|
|
|
|
x_quant = activation_quant(x).to(target_dtype) |
|
|
|
|
|
if self.bias is not None: |
|
|
b = self.bias.to(target_dtype) |
|
|
else: |
|
|
b = None |
|
|
|
|
|
return F.linear(x_quant, w_quant, b) |
|
|
|
|
|
def convert_to_bitnet(model, copy_weights=False): |
|
|
for name, module in model.named_children(): |
|
|
if isinstance(module, nn.Linear): |
|
|
bit_linear = BitLinear(module.in_features, module.out_features, module.bias is not None) |
|
|
if copy_weights: |
|
|
bit_linear.weight.data = module.weight.data.clone() |
|
|
if module.bias is not None: |
|
|
bit_linear.bias.data = module.bias.data.clone() |
|
|
setattr(model, name, bit_linear) |
|
|
else: |
|
|
convert_to_bitnet(module, copy_weights=copy_weights) |
|
|
|
|
|
class JobStatus: |
|
|
def __init__(self): |
|
|
self.id = str(uuid.uuid4()) |
|
|
self.status = "INITIALIZING" |
|
|
self.progress = 0.0 |
|
|
self.logs = [] |
|
|
self.result = None |
|
|
self.error = None |
|
|
self.created_at = datetime.now().strftime("%H:%M:%S") |
|
|
self.repo_url = None |
|
|
|
|
|
def add_log(self, message): |
|
|
timestamp = datetime.now().strftime("%H:%M:%S") |
|
|
self.logs.append(f"[{timestamp}] {message}") |
|
|
|
|
|
def set_progress(self, val, msg=None): |
|
|
self.progress = val |
|
|
if msg: |
|
|
self.add_log(msg) |
|
|
|
|
|
class CustomTrainerCallback(TrainerCallback): |
|
|
def __init__(self, job_id, hf_token, repo_id): |
|
|
self.job_id = job_id |
|
|
self.hf_token = hf_token |
|
|
self.repo_id = repo_id |
|
|
|
|
|
def on_step_end(self, args, state, control, **kwargs): |
|
|
if self.job_id in JOBS: |
|
|
job = JOBS[self.job_id] |
|
|
if state.max_steps > 0: |
|
|
prog = state.global_step / state.max_steps |
|
|
job.progress = 0.1 + (prog * 0.8) |
|
|
if state.global_step % 1 == 0: |
|
|
loss = state.log_history[-1].get('loss', 'N/A') if state.log_history else '...' |
|
|
job.add_log(f"Training Step {state.global_step}/{state.max_steps} | Loss: {loss}") |
|
|
return control |
|
|
|
|
|
def on_save(self, args, state, control, **kwargs): |
|
|
if self.job_id in JOBS: |
|
|
job = JOBS[self.job_id] |
|
|
step = state.global_step |
|
|
ckpt_name = f"checkpoint-{step}" |
|
|
ckpt_path = os.path.join(args.output_dir, ckpt_name) |
|
|
|
|
|
job.add_log(f"System: 100-Step Snapshot saved ({ckpt_name})") |
|
|
|
|
|
def _upload_bg(): |
|
|
try: |
|
|
upload_folder( |
|
|
folder_path=ckpt_path, |
|
|
path_in_repo=".", |
|
|
repo_id=self.repo_id, |
|
|
token=self.hf_token, |
|
|
commit_message=f"Live Checkpoint Step {step}" |
|
|
) |
|
|
job.add_log(f"Cloud: Synced Checkpoint {step} to Root") |
|
|
except: |
|
|
pass |
|
|
|
|
|
threading.Thread(target=_upload_bg, daemon=True).start() |
|
|
return control |
|
|
|
|
|
@spaces.GPU(duration=300) |
|
|
def background_train_task(job_id, hf_token, model_name, new_repo_name, |
|
|
train_steps, learning_rate, batch_size, datasets_text, |
|
|
reasoning_mode, c_conf, c_tok, c_gen): |
|
|
|
|
|
job = JOBS[job_id] |
|
|
job.status = "RUNNING" |
|
|
job.add_log("System: initializing BitNet Scratch Protocol...") |
|
|
|
|
|
try: |
|
|
if not hf_token.startswith("hf_"): |
|
|
raise ValueError("Invalid Token") |
|
|
|
|
|
os.environ["WANDB_DISABLED"] = "true" |
|
|
os.environ["HF_TOKEN"] = hf_token |
|
|
os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true" |
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "true" |
|
|
|
|
|
login(token=hf_token) |
|
|
try: |
|
|
username = whoami()["name"] |
|
|
full_repo_id = f"{username}/{new_repo_name}" |
|
|
create_repo(full_repo_id, token=hf_token, exist_ok=True) |
|
|
job.add_log(f"Auth: Verified {username} -> {full_repo_id}") |
|
|
except: |
|
|
raise Exception("Auth Failed") |
|
|
|
|
|
if not hasattr(torch, 'xla'): |
|
|
class DummyXLA: |
|
|
def __getattr__(self, name): |
|
|
return lambda *args, **kwargs: None |
|
|
torch.xla = DummyXLA() |
|
|
|
|
|
raw_items = datasets_text.replace('\n', ',').split(',') |
|
|
dataset_list = [item.strip() for item in raw_items if item.strip()] |
|
|
|
|
|
if reasoning_mode: |
|
|
job.add_log("Config: Reasoning Injection Active") |
|
|
dataset_list.extend(["gsm8k", "openai/gsm8k"]) |
|
|
|
|
|
def load_single(ds_name, cfg): |
|
|
try: |
|
|
ds = load_dataset(ds_name, cfg if cfg else "main", split="train", streaming=True, trust_remote_code=False) |
|
|
try: |
|
|
next(iter(ds)) |
|
|
return ds |
|
|
except: |
|
|
return None |
|
|
except: |
|
|
return None |
|
|
|
|
|
streams = [] |
|
|
job.set_progress(0.05, "Data: Parallel Stream Connect...") |
|
|
|
|
|
cpu_count = multiprocessing.cpu_count() |
|
|
with ThreadPoolExecutor(max_workers=cpu_count * 2) as executor: |
|
|
futures = [] |
|
|
for ds_name in dataset_list: |
|
|
futures.append(executor.submit(load_single, ds_name, None)) |
|
|
|
|
|
for future in as_completed(futures): |
|
|
res = future.result() |
|
|
if res: |
|
|
streams.append(res) |
|
|
|
|
|
if not streams: |
|
|
raise Exception("No Data Sources") |
|
|
|
|
|
job.set_progress(0.1, f"Data: {len(streams)} Streams Linked") |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="left", add_eos_token=True, add_bos_token=True) |
|
|
if tokenizer.pad_token is None: |
|
|
tokenizer.pad_token = tokenizer.eos_token |
|
|
|
|
|
def process_stream_generator(): |
|
|
iterator = chain.from_iterable(streams) |
|
|
batch_buffer = [] |
|
|
|
|
|
for item in iterator: |
|
|
try: |
|
|
text = str(item.get("text", item.get("content", str(item)))) |
|
|
if len(text) < 5: continue |
|
|
batch_buffer.append(text) |
|
|
|
|
|
if len(batch_buffer) >= 100: |
|
|
encoded_batch = tokenizer(batch_buffer, truncation=True, max_length=2048, padding=False) |
|
|
for input_ids in encoded_batch["input_ids"]: |
|
|
yield {"input_ids": input_ids} |
|
|
batch_buffer = [] |
|
|
except: |
|
|
continue |
|
|
|
|
|
job.set_progress(0.15, "Model: Initializing Architecture & Converting to BitNet...") |
|
|
|
|
|
torch.cuda.empty_cache() |
|
|
gc.collect() |
|
|
|
|
|
config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) |
|
|
|
|
|
original_model = AutoModelForCausalLM.from_config( |
|
|
config, |
|
|
trust_remote_code=True, |
|
|
) |
|
|
|
|
|
convert_to_bitnet(original_model, copy_weights=False) |
|
|
|
|
|
model_size = sum(t.numel() for t in original_model.parameters()) |
|
|
job.add_log(f"Model Size: {model_size/1000**2:.1f}M Parameters (1.58-bit)") |
|
|
|
|
|
output_dir = f"checkpoints/{job_id}" |
|
|
|
|
|
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False) |
|
|
|
|
|
training_args = TrainingArguments( |
|
|
output_dir=output_dir, |
|
|
per_device_train_batch_size=int(batch_size), |
|
|
gradient_accumulation_steps=4, |
|
|
max_steps=int(train_steps), |
|
|
learning_rate=learning_rate, |
|
|
optim="adamw_torch_fused" if torch.cuda.is_available() else "adamw_torch", |
|
|
logging_steps=1, |
|
|
save_strategy="steps", |
|
|
save_steps=100, |
|
|
save_total_limit=1, |
|
|
report_to="none", |
|
|
fp16=True if torch.cuda.is_available() else False, |
|
|
disable_tqdm=True, |
|
|
dataloader_num_workers=4, |
|
|
dataloader_pin_memory=True, |
|
|
gradient_checkpointing=True, |
|
|
torch_compile=False, |
|
|
lr_scheduler_type="cosine", |
|
|
warmup_ratio=0.1 |
|
|
) |
|
|
|
|
|
dataset_iterable = IterableDataset.from_generator(process_stream_generator) |
|
|
|
|
|
trainer = Trainer( |
|
|
model=original_model, |
|
|
tokenizer=tokenizer, |
|
|
train_dataset=dataset_iterable, |
|
|
args=training_args, |
|
|
data_collator=data_collator, |
|
|
callbacks=[CustomTrainerCallback(job_id, hf_token, full_repo_id)] |
|
|
) |
|
|
|
|
|
job.set_progress(0.2, "Training: BitNet Gradient Descent Initiated...") |
|
|
trainer.train() |
|
|
trainer.save_model(output_dir) |
|
|
tokenizer.save_pretrained(output_dir) |
|
|
|
|
|
job.set_progress(0.9, "Processing: Finalizing Artifacts...") |
|
|
del original_model |
|
|
torch.cuda.empty_cache() |
|
|
gc.collect() |
|
|
|
|
|
def inject_json(content, fname): |
|
|
if content and content.strip(): |
|
|
try: |
|
|
data = json.loads(content) |
|
|
file_path = os.path.join(output_dir, fname) |
|
|
|
|
|
if os.path.exists(file_path): |
|
|
with open(file_path, 'r', encoding='utf-8') as f: |
|
|
try: |
|
|
existing_data = json.load(f) |
|
|
existing_data.update(data) |
|
|
data = existing_data |
|
|
except: |
|
|
pass |
|
|
|
|
|
with open(file_path, 'w', encoding='utf-8') as f: |
|
|
json.dump(data, f, indent=2) |
|
|
job.add_log(f"Config: Overwritten {fname} with user settings") |
|
|
except: |
|
|
pass |
|
|
|
|
|
inject_json(c_conf, "config.json") |
|
|
inject_json(c_tok, "tokenizer_config.json") |
|
|
inject_json(c_gen, "generation_config.json") |
|
|
|
|
|
job.set_progress(0.95, "Network: Uploading Final BitNet Model...") |
|
|
|
|
|
upload_folder( |
|
|
folder_path=output_dir, |
|
|
path_in_repo=".", |
|
|
repo_id=full_repo_id, |
|
|
token=hf_token, |
|
|
commit_message="BitNet Scratch Trained Model" |
|
|
) |
|
|
|
|
|
job.repo_url = f"https://huggingface.co/{full_repo_id}" |
|
|
job.status = "COMPLETED" |
|
|
job.set_progress(1.0, "System: Operation Finalized") |
|
|
|
|
|
except Exception as e: |
|
|
job.status = "FAILED" |
|
|
job.error = str(e) |
|
|
job.add_log(f"FATAL ERROR: {str(e)}") |
|
|
torch.cuda.empty_cache() |
|
|
|
|
|
def start_training_wrapper(hf_token, model_name, new_repo_name, |
|
|
train_steps, learning_rate, batch_size, datasets_text, |
|
|
reasoning_mode, c_conf, c_tok, c_gen): |
|
|
|
|
|
if not hf_token or not model_name: |
|
|
return None, gr.update(selected="launch_tab") |
|
|
|
|
|
new_job = JobStatus() |
|
|
JOBS[new_job.id] = new_job |
|
|
|
|
|
thread = threading.Thread( |
|
|
target=background_train_task, |
|
|
args=(new_job.id, hf_token, model_name, new_repo_name, |
|
|
train_steps, learning_rate, batch_size, datasets_text, reasoning_mode, c_conf, c_tok, c_gen) |
|
|
) |
|
|
thread.daemon = True |
|
|
thread.start() |
|
|
|
|
|
return new_job.id, gr.update(selected="monitor_tab") |
|
|
|
|
|
def get_job_update(job_id): |
|
|
if not job_id: |
|
|
return "Waiting...", "", 0, "", gr.update(visible=False) |
|
|
|
|
|
if job_id not in JOBS: |
|
|
return "Not Found", "", 0, "", gr.update(visible=False) |
|
|
|
|
|
job = JOBS[job_id] |
|
|
|
|
|
log_text = "\n".join(job.logs) |
|
|
|
|
|
result_comp = gr.update(visible=False) |
|
|
if job.status == "COMPLETED" and job.repo_url: |
|
|
result_comp = gr.update(visible=True, value=f"✅ Full Model Published: {job.repo_url}") |
|
|
|
|
|
return job.status, job.created_at, job.progress, log_text, result_comp |
|
|
|
|
|
def load_from_url(request: gr.Request): |
|
|
try: |
|
|
params = request.query_params |
|
|
job_id = params.get("job_id") |
|
|
if job_id: |
|
|
return gr.update(selected="monitor_tab"), job_id |
|
|
except: |
|
|
pass |
|
|
return gr.update(selected="launch_tab"), "" |
|
|
|
|
|
with gr.Blocks(title="Nucleus Enterprise") as demo: |
|
|
with gr.Column(): |
|
|
gr.Markdown("# ⚛️ NUCLEUS ENTERPRISE") |
|
|
gr.Markdown("Autonomous LLM Foundry | V10.0 BitNet Edition") |
|
|
|
|
|
with gr.Tabs() as main_tabs: |
|
|
with gr.TabItem("🚀 LAUNCHPAD", id="launch_tab"): |
|
|
with gr.Row(): |
|
|
with gr.Column(scale=2): |
|
|
with gr.Row(): |
|
|
hf_token = gr.Textbox(label="HuggingFace Token", type="password", value=os.getenv("HF_TOKEN", "")) |
|
|
model_name = gr.Textbox(label="Architecture Config Source", value="Qwen/Qwen2.5-0.5B") |
|
|
|
|
|
repo_name = gr.Textbox(label="Output Repository", value="nucleus-bitnet-v1") |
|
|
datasets = gr.Textbox(label="Datasets (CSV)", value="Salesforce/fineweb_deduplicated", lines=3) |
|
|
|
|
|
reasoning = gr.Checkbox(label="Inject Reasoning (CoT/Math)", value=False) |
|
|
|
|
|
with gr.Column(scale=1): |
|
|
steps = gr.Number(label="Steps", value=100) |
|
|
lr = gr.Number(label="Learning Rate", value=1e-4) |
|
|
batch = gr.Number(label="Batch Size", value=1) |
|
|
|
|
|
with gr.Accordion("Advanced Config", open=False): |
|
|
c_conf = gr.Code(label="config.json", language="json") |
|
|
c_tok = gr.Code(label="tokenizer_config.json", language="json") |
|
|
c_gen = gr.Code(label="generation_config.json", language="json") |
|
|
|
|
|
btn_launch = gr.Button("INITIALIZE BITNET TRAINING", variant="primary", size="lg") |
|
|
|
|
|
with gr.TabItem("📡 TELEMETRY", id="monitor_tab"): |
|
|
with gr.Row(): |
|
|
job_id_input = gr.Textbox(label="Active Job ID", interactive=True) |
|
|
btn_refresh = gr.Button("Refresh Stream") |
|
|
|
|
|
with gr.Row(): |
|
|
status_out = gr.Textbox(label="Status", interactive=False) |
|
|
time_out = gr.Textbox(label="Start Time", interactive=False) |
|
|
progress_out = gr.Slider(label="Progress", minimum=0, maximum=1) |
|
|
|
|
|
final_link = gr.Markdown(visible=False) |
|
|
logs_out = gr.Code(label="Real-time Kernel Logs", language="shell", interactive=False, lines=15) |
|
|
|
|
|
timer = gr.Timer(2000, active=False) |
|
|
|
|
|
demo.load(load_from_url, None, [main_tabs, job_id_input]).then(lambda: gr.Timer(active=True), None, timer) |
|
|
|
|
|
btn_launch.click( |
|
|
start_training_wrapper, |
|
|
inputs=[hf_token, model_name, repo_name, steps, lr, batch, datasets, reasoning, c_conf, c_tok, c_gen], |
|
|
outputs=[job_id_input, main_tabs] |
|
|
).then( |
|
|
None, [job_id_input], None, |
|
|
js="(id) => { if (id) { const url = new URL(window.location); url.searchParams.set('job_id', id); window.history.pushState({}, '', url); } return id; }" |
|
|
).then( |
|
|
lambda: gr.Timer(active=True), None, timer |
|
|
) |
|
|
|
|
|
btn_refresh.click(get_job_update, job_id_input, [status_out, time_out, progress_out, logs_out, final_link]) |
|
|
timer.tick(get_job_update, job_id_input, [status_out, time_out, progress_out, logs_out, final_link]) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
demo.launch(ssr_mode=False) |
